<commit_before>#!/usr/bin/env python
'''
PEXPECT LICENSE

    This license is approved by the OSI and FSF as GPL-compatible.
        http://opensource.org/licenses/isc-license.txt

    Copyright (c) 2012, Noah Spurrier <noah@noah.org>
    PERMISSION TO USE, COPY, MODIFY, AND/OR DISTRIBUTE THIS SOFTWARE FOR ANY
    PURPOSE WITH OR WITHOUT FEE IS HEREBY GRANTED, PROVIDED THAT THE ABOVE
    COPYRIGHT NOTICE AND THIS PERMISSION NOTICE APPEAR IN ALL COPIES.
    THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
    WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
    MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
    ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
    WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
    ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
    OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

'''
from __future__ import with_statement  # bring 'with' stmt to py25
import pexpect
import unittest
import PexpectTestCase
import os
import re
from pexpect import six

# the program cat(1) may display ^D\x08\x08 when \x04 (EOF, Ctrl-D) is sent
_CAT_EOF = six.b('^D\x08\x08')

class TestCaseMisc(PexpectTestCase.PexpectTestCase):

    def test_isatty (self):
        child = pexpect.spawn('cat')
        assert child.isatty(), "Not returning True. Should always be True."

    def test_read (self):
        child = pexpect.spawn('cat')
        child.sendline ("abc")
        child.sendeof()
        self.assertEqual(child.read(0), six.b(''))
        self.assertEqual(child.read(1), six.b('a'))
        self.assertEqual(child.read(1), six.b('b'))
        self.assertEqual(child.read(1), six.b('c'))
        self.assertEqual(child.read(2), six.b('\r\n'))
        remaining = child.read().replace(_CAT_EOF, six.b(''))
        self.assertEqual(remaining, six.b('abc\r\n'))

    def test_readline (self):
        '''See the note in test_readlines() for an explaination as to why
        I allow line3 and line4 to return multiple patterns.
        Basically, this is done to handle a valid condition on slow systems.
        '''
        child = pexpect.spawn('cat')
        child.sendline ("abc")
        child.sendline ("123")
        child.sendeof()
        line1 = child.readline(0)
        line2 = child.readline()
        line3 = child.readline(2)
        line4 = child.readline(1)
        line5 = child.readline()
        assert child.isalive() is False, child.isalive()
        assert child.exitstatus == 0, child.exitstatus
        self.assertEqual(line1, six.b(''))
        self.assertEqual(line2, six.b('abc\r\n'))
        assert (line3 == six.b('abc\r\n') or line3 == '123\r\n'), \
            "readline(2) did not return 'abc\\r\\n'. Returned: " + repr(line3)
        assert (line4 == six.b('123\r\n') or line4 == 'abc\r\n'), \
            "readline(1) did not return '123\\r\\n'. Returned: " + repr(line4)
        self.assertEqual(line5, six.b('123\r\n'))

    def test_iter (self):
        '''See the note in test_readlines() for an explaination as to why
        I allow line3 and line4 to return multiple patterns.
        Basically, this is done to handle a valid condition on slow systems.
        '''
        child = pexpect.spawn('cat')
        child.sendline ("abc")
        child.sendline ("123")
        child.sendeof()
        # Don't use ''.join() because we want to test the ITERATOR.
        page = six.b('')
        for line in child:
            page += line
        page = page.replace(_CAT_EOF, six.b(''))
        assert (page == six.b('abc\r\nabc\r\n123\r\n123\r\n') or
                page == six.b('abc\r\n123\r\nabc\r\n123\r\n')) , \
               "iterator did not work. page=%s"%repr(page)

    def test_readlines(self):
        '''Note that on some slow or heavily loaded systems that the lines
        coming back from 'cat' may come even after the EOF.
        We except to see two copies of the lines we send 'cat'.
        The first line is the TTY echo, the second line is from 'cat'.
        Usually 'cat' will respond with 'abc' before we have a chance to
        send the second line, '123'. If this does not happen then the
        lines may appear out of order. This is technically not an error.
        That's just the nature of asynchronous communication.
        This is why the assert will allow either of the two possible
        patterns to be returned by lineslined(). The (lame) alternative is
        to put a long sleep between the two sendline() calls, but then
        I have to make assumptions about how fast 'cat' can reply.
        '''
        child = pexpect.spawn('cat')
        child.sendline ("abc")
        child.sendline ("123")
        child.sendeof()
        assert child.isalive() is False, child.isalive
        assert child.exitstatus == 0, child.exitstatus
        page = six.b('').join(child.readlines()).replace(_CAT_EOF, six.b(''))
        assert (page == six.b('abc\r\nabc\r\n123\r\n123\r\n') or
                page == six.b('abc\r\n123\r\nabc\r\n123\r\n')), \
               "readlines() did not work. page=%s"%repr(page)

    def test_write (self):
        child = pexpect.spawn('cat')
        child.write('a')
        child.write('\r')
        self.assertEqual(child.readline(), six.b('a\r\n'))

    def test_writelines (self):
        child = pexpect.spawn('cat')
        child.writelines(['abc','123','xyz','\r'])
        child.sendeof()
        line = child.readline()
        assert line == six.b('abc123xyz\r\n'), "writelines() did not work. line=%s"%repr(line)

    def test_eof(self):
        child = pexpect.spawn('cat')
        child.sendeof()
        try:
            child.expect ('the unexpected')
        except:
            pass
        assert child.eof(), "child.eof() did not return True"

    def test_terminate(self):
        child = pexpect.spawn('cat')
        child.terminate(force=1)
        assert child.terminated, "child.terminated is not True"

    def test_bad_child_pid(self):
        child = pexpect.spawn('cat')
        child.terminate(force=1)
        child.terminated = 0 # Force invalid state to test code
        try:
            child.isalive()
        except pexpect.ExceptionPexpect:
            pass
        else:
            self.fail ("child.isalive() should have raised a pexpect.ExceptionPexpect")
        child.terminated = 1 # Force back to valid state so __del__ won't complain
    def test_bad_arguments (self):
        '''This tests that we get a graceful error when passing bad arguments.'''
        try:
            p = pexpect.spawn(1)
        except pexpect.ExceptionPexpect:
            pass
        else:
            self.fail ("pexpect.spawn(1) should have raised a pexpect.ExceptionPexpect.")
        try:
            p = pexpect.spawn('ls', '-la') # should really use pexpect.spawn('ls', ['-ls'])
        except TypeError:
            pass
        else:
            self.fail ("pexpect.spawn('ls', '-la') should have raised a TypeError.")
        try:
            p = pexpect.spawn('cat')
            p.close()
            p.read_nonblocking(size=1, timeout=3)
        except ValueError:
            pass
        else:
            self.fail ("read_nonblocking on closed spawn object should have raised a ValueError.")
    def test_isalive(self):
        child = pexpect.spawn('cat')
        assert child.isalive(), "child.isalive() did not return True"
        child.sendeof()
        child.expect(pexpect.EOF)
        assert not child.isalive(), "child.isalive() did not return False"
    def test_bad_type_in_expect(self):
        child = pexpect.spawn('cat')
        try:
            child.expect({}) # We don't support dicts yet. Should give TypeError
        except TypeError:
            pass
        else:
            self.fail ("child.expect({}) should have raised a TypeError")
    def test_winsize(self):
        child = pexpect.spawn('cat')
        child.setwinsize(10,13)
        assert child.getwinsize()==(10,13), "getwinsize() did not return (10,13)"

    def test_env(self):
        default = pexpect.run('env')
        userenv = pexpect.run('env', env={'foo':'pexpect'})
        assert default!=userenv, "'default' and 'userenv' should be different"
        assert six.b('foo') in userenv and six.b('pexpect') in userenv, "'foo' and 'pexpect' should be in 'userenv'"

    def test_cwd (self): # This assumes 'pwd' and '/tmp' exist on this platform.
        default = pexpect.run('pwd')
        tmpdir =  pexpect.run('pwd', cwd='/tmp')
        assert default!=tmpdir, "'default' and 'tmpdir' should be different"
        assert (six.b('tmp') in tmpdir), "'tmp' should be returned by 'pwd' command"

    def test_which (self):
        p = os.defpath
        ep = os.environ['PATH']
        os.defpath = ":/tmp"
        os.environ['PATH'] = ":/tmp"
        wp = pexpect.which ("ticker.py")
        assert wp == 'ticker.py', "Should return a string. Returned %s" % wp
        os.defpath = "/tmp"
        os.environ['PATH'] = "/tmp"
        wp = pexpect.which ("ticker.py")
        assert wp == None, "Executable should not be found. Returned %s" % wp
        os.defpath = p
        os.environ['PATH'] = ep

    def test_searcher_re (self):
        ss = pexpect.searcher_re ([re.compile('this'),re.compile('that'),re.compile('and'),re.compile('the'),re.compile('other')])
        assert ss.__str__() == 'searcher_re:\n    0: re.compile("this")\n    1: re.compile("that")\n    2: re.compile("and")\n    3: re.compile("the")\n    4: re.compile("other")'
        ss = pexpect.searcher_re ([pexpect.TIMEOUT,re.compile('this'),re.compile('that'),re.compile('and'),pexpect.EOF,re.compile('other')])
        assert ss.__str__() == 'searcher_re:\n    0: TIMEOUT\n    1: re.compile("this")\n    2: re.compile("that")\n    3: re.compile("and")\n    4: EOF\n    5: re.compile("other")'

    def test_searcher_string (self):
        ss = pexpect.searcher_string (['this','that','and','the','other'])
        assert ss.__str__() == 'searcher_string:\n    0: "this"\n    1: "that"\n    2: "and"\n    3: "the"\n    4: "other"', repr(ss.__str__())
        ss = pexpect.searcher_string (['this',pexpect.EOF,'that','and','the','other',pexpect.TIMEOUT])
        assert ss.__str__() == 'searcher_string:\n    0: "this"\n    1: EOF\n    2: "that"\n    3: "and"\n    4: "the"\n    5: "other"\n    6: TIMEOUT'

if __name__ == '__main__':
    unittest.main()

suite = unittest.makeSuite(TestCaseMisc,'test')

<commit_msg>display error fix and brevity of childexit assrts<commit_after>#!/usr/bin/env python
'''
PEXPECT LICENSE

    This license is approved by the OSI and FSF as GPL-compatible.
        http://opensource.org/licenses/isc-license.txt

    Copyright (c) 2012, Noah Spurrier <noah@noah.org>
    PERMISSION TO USE, COPY, MODIFY, AND/OR DISTRIBUTE THIS SOFTWARE FOR ANY
    PURPOSE WITH OR WITHOUT FEE IS HEREBY GRANTED, PROVIDED THAT THE ABOVE
    COPYRIGHT NOTICE AND THIS PERMISSION NOTICE APPEAR IN ALL COPIES.
    THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
    WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
    MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
    ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
    WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
    ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
    OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

'''
from __future__ import with_statement  # bring 'with' stmt to py25
import pexpect
import unittest
import PexpectTestCase
import os
import re
from pexpect import six

# the program cat(1) may display ^D\x08\x08 when \x04 (EOF, Ctrl-D) is sent
_CAT_EOF = six.b('^D\x08\x08')

class TestCaseMisc(PexpectTestCase.PexpectTestCase):

    def test_isatty (self):
        child = pexpect.spawn('cat')
        assert child.isatty(), "Not returning True. Should always be True."

    def test_read (self):
        child = pexpect.spawn('cat')
        child.sendline ("abc")
        child.sendeof()
        self.assertEqual(child.read(0), six.b(''))
        self.assertEqual(child.read(1), six.b('a'))
        self.assertEqual(child.read(1), six.b('b'))
        self.assertEqual(child.read(1), six.b('c'))
        self.assertEqual(child.read(2), six.b('\r\n'))
        remaining = child.read().replace(_CAT_EOF, six.b(''))
        self.assertEqual(remaining, six.b('abc\r\n'))

    def test_readline (self):
        '''See the note in test_readlines() for an explaination as to why
        I allow line3 and line4 to return multiple patterns.
        Basically, this is done to handle a valid condition on slow systems.
        '''
        child = pexpect.spawn('cat')
        child.sendline ("abc")
        child.sendline ("123")
        child.sendeof()
        line1 = child.readline(0)
        line2 = child.readline()
        line3 = child.readline(2)
        line4 = child.readline(1)
        line5 = child.readline()
        assert not child.isalive(), child.isalive()
        assert child.exitstatus == 0, child.exitstatus
        self.assertEqual(line1, six.b(''))
        self.assertEqual(line2, six.b('abc\r\n'))
        assert (line3 == six.b('abc\r\n') or line3 == '123\r\n'), \
            "readline(2) did not return 'abc\\r\\n'. Returned: " + repr(line3)
        assert (line4 == six.b('123\r\n') or line4 == 'abc\r\n'), \
            "readline(1) did not return '123\\r\\n'. Returned: " + repr(line4)
        self.assertEqual(line5, six.b('123\r\n'))

    def test_iter (self):
        '''See the note in test_readlines() for an explaination as to why
        I allow line3 and line4 to return multiple patterns.
        Basically, this is done to handle a valid condition on slow systems.
        '''
        child = pexpect.spawn('cat')
        child.sendline ("abc")
        child.sendline ("123")
        child.sendeof()
        # Don't use ''.join() because we want to test the ITERATOR.
        page = six.b('')
        for line in child:
            page += line
        page = page.replace(_CAT_EOF, six.b(''))
        assert (page == six.b('abc\r\nabc\r\n123\r\n123\r\n') or
                page == six.b('abc\r\n123\r\nabc\r\n123\r\n')) , \
               "iterator did not work. page=%s"%repr(page)

    def test_readlines(self):
        '''Note that on some slow or heavily loaded systems that the lines
        coming back from 'cat' may come even after the EOF.
        We except to see two copies of the lines we send 'cat'.
        The first line is the TTY echo, the second line is from 'cat'.
        Usually 'cat' will respond with 'abc' before we have a chance to
        send the second line, '123'. If this does not happen then the
        lines may appear out of order. This is technically not an error.
        That's just the nature of asynchronous communication.
        This is why the assert will allow either of the two possible
        patterns to be returned by lineslined(). The (lame) alternative is
        to put a long sleep between the two sendline() calls, but then
        I have to make assumptions about how fast 'cat' can reply.
        '''
        child = pexpect.spawn('cat')
        child.sendline ("abc")
        child.sendline ("123")
        child.sendeof()
        assert not child.isalive(), child.isalive()
        assert child.exitstatus == 0, child.exitstatus
        page = six.b('').join(child.readlines()).replace(_CAT_EOF, six.b(''))
        assert (page == six.b('abc\r\nabc\r\n123\r\n123\r\n') or
                page == six.b('abc\r\n123\r\nabc\r\n123\r\n')), \
               "readlines() did not work. page=%s"%repr(page)

    def test_write (self):
        child = pexpect.spawn('cat')
        child.write('a')
        child.write('\r')
        self.assertEqual(child.readline(), six.b('a\r\n'))

    def test_writelines (self):
        child = pexpect.spawn('cat')
        child.writelines(['abc','123','xyz','\r'])
        child.sendeof()
        line = child.readline()
        assert line == six.b('abc123xyz\r\n'), "writelines() did not work. line=%s"%repr(line)

    def test_eof(self):
        child = pexpect.spawn('cat')
        child.sendeof()
        try:
            child.expect ('the unexpected')
        except:
            pass
        assert child.eof(), "child.eof() did not return True"

    def test_terminate(self):
        child = pexpect.spawn('cat')
        child.terminate(force=1)
        assert child.terminated, "child.terminated is not True"

    def test_bad_child_pid(self):
        child = pexpect.spawn('cat')
        child.terminate(force=1)
        child.terminated = 0 # Force invalid state to test code
        try:
            child.isalive()
        except pexpect.ExceptionPexpect:
            pass
        else:
            self.fail ("child.isalive() should have raised a pexpect.ExceptionPexpect")
        child.terminated = 1 # Force back to valid state so __del__ won't complain
    def test_bad_arguments (self):
        '''This tests that we get a graceful error when passing bad arguments.'''
        try:
            p = pexpect.spawn(1)
        except pexpect.ExceptionPexpect:
            pass
        else:
            self.fail ("pexpect.spawn(1) should have raised a pexpect.ExceptionPexpect.")
        try:
            p = pexpect.spawn('ls', '-la') # should really use pexpect.spawn('ls', ['-ls'])
        except TypeError:
            pass
        else:
            self.fail ("pexpect.spawn('ls', '-la') should have raised a TypeError.")
        try:
            p = pexpect.spawn('cat')
            p.close()
            p.read_nonblocking(size=1, timeout=3)
        except ValueError:
            pass
        else:
            self.fail ("read_nonblocking on closed spawn object should have raised a ValueError.")
    def test_isalive(self):
        child = pexpect.spawn('cat')
        assert child.isalive(), child.isalive()
        child.sendeof()
        child.expect(pexpect.EOF)
        assert not child.isalive(), child.isalive()
    def test_bad_type_in_expect(self):
        child = pexpect.spawn('cat')
        try:
            child.expect({}) # We don't support dicts yet. Should give TypeError
        except TypeError:
            pass
        else:
            self.fail ("child.expect({}) should have raised a TypeError")
    def test_winsize(self):
        child = pexpect.spawn('cat')
        child.setwinsize(10,13)
        assert child.getwinsize()==(10,13), "getwinsize() did not return (10,13)"

    def test_env(self):
        default = pexpect.run('env')
        userenv = pexpect.run('env', env={'foo':'pexpect'})
        assert default!=userenv, "'default' and 'userenv' should be different"
        assert six.b('foo') in userenv and six.b('pexpect') in userenv, "'foo' and 'pexpect' should be in 'userenv'"

    def test_cwd (self): # This assumes 'pwd' and '/tmp' exist on this platform.
        default = pexpect.run('pwd')
        tmpdir =  pexpect.run('pwd', cwd='/tmp')
        assert default!=tmpdir, "'default' and 'tmpdir' should be different"
        assert (six.b('tmp') in tmpdir), "'tmp' should be returned by 'pwd' command"

    def test_which (self):
        p = os.defpath
        ep = os.environ['PATH']
        os.defpath = ":/tmp"
        os.environ['PATH'] = ":/tmp"
        wp = pexpect.which ("ticker.py")
        assert wp == 'ticker.py', "Should return a string. Returned %s" % wp
        os.defpath = "/tmp"
        os.environ['PATH'] = "/tmp"
        wp = pexpect.which ("ticker.py")
        assert wp == None, "Executable should not be found. Returned %s" % wp
        os.defpath = p
        os.environ['PATH'] = ep

    def test_searcher_re (self):
        ss = pexpect.searcher_re ([re.compile('this'),re.compile('that'),re.compile('and'),re.compile('the'),re.compile('other')])
        assert ss.__str__() == 'searcher_re:\n    0: re.compile("this")\n    1: re.compile("that")\n    2: re.compile("and")\n    3: re.compile("the")\n    4: re.compile("other")'
        ss = pexpect.searcher_re ([pexpect.TIMEOUT,re.compile('this'),re.compile('that'),re.compile('and'),pexpect.EOF,re.compile('other')])
        assert ss.__str__() == 'searcher_re:\n    0: TIMEOUT\n    1: re.compile("this")\n    2: re.compile("that")\n    3: re.compile("and")\n    4: EOF\n    5: re.compile("other")'

    def test_searcher_string (self):
        ss = pexpect.searcher_string (['this','that','and','the','other'])
        assert ss.__str__() == 'searcher_string:\n    0: "this"\n    1: "that"\n    2: "and"\n    3: "the"\n    4: "other"', repr(ss.__str__())
        ss = pexpect.searcher_string (['this',pexpect.EOF,'that','and','the','other',pexpect.TIMEOUT])
        assert ss.__str__() == 'searcher_string:\n    0: "this"\n    1: EOF\n    2: "that"\n    3: "and"\n    4: "the"\n    5: "other"\n    6: TIMEOUT'

if __name__ == '__main__':
    unittest.main()

suite = unittest.makeSuite(TestCaseMisc,'test')

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

"""This version is also read by setup.py."""

from pygments.token import Token
Token.Name.XojoType

__version__ = '0.0.0'
<commit_msg>Satisfy pylint.<commit_after># -*- coding: utf-8 -*-

"""Pygments syntax highlighting for Xojo."""

from pygments.token import Token
# the point of the next statement is that accessing the attribute creates the token type,
# as explained in the pygments token documentation.
#pylint:disable=pointless-statement
Token.Name.XojoType

# __version is also read by setup.py.
__version__ = '0.0.0'
<|endoftext|>


<commit_before>"""
Tests for plotting operations
"""

import os
import unittest
from unittest import TestCase

import numpy as np
import xarray as xr

from cate.ops.plot import plot, plot_map

@unittest.skipIf(condition=os.environ.get('CATE_DISABLE_PLOT_TESTS', None),
                 reason="skipped if CATE_DISABLE_PLOT_TESTS=1")
class TestPlotMap(TestCase):
    """
    Test plot_map() function
    """
    def test_plot_map(self):
        # Test the nominal functionality. This doesn't check that the plot is what is expected,
        # rather, it simply tests if it seems to have been created
        dataset = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.random.rand(5, 10, 2)),
            'second': (['lat', 'lon', 'time'], np.random.rand(5, 10, 2)),
            'lat': np.linspace(-89.5, 89.5, 5),
            'lon': np.linspace(-179.5, 179.5, 10)})

        plot_map(dataset, file='remove_me.png')
        self.assertTrue(os.path.isfile('remove_me.png'))
        os.remove('remove_me.png')

        # Test if an error is raised when an unsupported format is passed
        with self.assertRaises(ValueError):
            plot_map(dataset, file='remove_me.pgm')
        self.assertFalse(os.path.isfile('remove_me.pgm'))

        # Test if extents can be used
        plot_map(dataset,
                 var='second',
                 region='-40.0, -20.0, 50.0, 60.0',
                 file='remove_me.pdf')
        self.assertTrue(os.path.isfile('remove_me.pdf'))
        os.remove('remove_me.pdf')

    def test_plot_map_exceptions(self):
        # Test if the corner cases are detected without creating a plot for it.

        # Test value error is raised when passing an unexpected dataset type
        with self.assertRaises(NotImplementedError):
            plot_map([1, 2, 4], file='remove_me.jpeg')
        self.assertFalse(os.path.isfile('remove_me.jpg'))

        # Test the extensions bound checking
        dataset = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.random.rand(2, 4, 1)),
            'lat': np.linspace(-89.5, 89.5, 2),
            'lon': np.linspace(-179.5, 179.5, 4)})

        with self.assertRaises(ValueError):
            region = '-40.0, -95.0, 50.0, 60.0',
            plot_map(dataset, region=region)

        with self.assertRaises(ValueError):
            region = '-40.0, -20.0, 50.0, 95.0',
            plot_map(dataset, region=region)

        with self.assertRaises(ValueError):
            region = '-181.0, -20.0, 50.0, 60.0',
            plot_map(dataset, region=region)

        with self.assertRaises(ValueError):
            region = '-40.0, -20.0, 181.0, 60.0',
            plot_map(dataset, region=region)

        with self.assertRaises(ValueError):
            region = '-40.0, -20.0, 50.0, -25.0',
            plot_map(dataset, region=region)

        with self.assertRaises(ValueError):
            region = '-20.0, -20.0, -25.0, 60.0',
            plot_map(dataset, region=region)

    def test_registered(self):
        """
        Test nominal execution of the function as a registered operation.
        """
        pass

@unittest.skipIf(condition=os.environ.get('CATE_DISABLE_PLOT_TESTS', None),
                 reason="skipped if CATE_DISABLE_PLOT_TESTS=1")
class TestPlot(TestCase):
    """
    Test plot() function
    """
    def test_plot(self):
        # Test plot
        dataset = xr.Dataset({
            'first': np.random.rand(10)})

        plot(dataset, 'first', file='remove_me.jpg')
        self.assertTrue(os.path.isfile('remove_me.jpg'))
        os.remove('remove_me.jpg')

    def test_registered(self):
        """
        Test nominal execution of the function as a registered operation.
        """
        pass

@unittest.skipIf(condition=os.environ.get('CATE_DISABLE_PLOT_TESTS', None),
                 reason="skipped if CATE_DISABLE_PLOT_TESTS=1")
class TestPlotDataFrame(TestCase):
    """
    Test plot_dataframe() function.
    """
    def test_nominal(self):
        """
        Test nominal execution
        """
        pass

    def test_registered(self):
        """
        Test the method when run as a registered operation
        """
        pass
<commit_msg>Use saner temp file management<commit_after>"""
Tests for plotting operations
"""

import os
import sys
import unittest
from unittest import TestCase

import numpy as np
import xarray as xr
import pandas as pd
import tempfile
import shutil
from contextlib import contextmanager
import itertools

from cate.core.op import OP_REGISTRY
from cate.util.misc import object_to_qualified_name

from cate.ops.plot import plot, plot_map, plot_dataframe

_counter = itertools.count()
ON_WIN = sys.platform == 'win32'

@contextmanager
def create_tmp_file(name, ext):
    tmp_dir = tempfile.mkdtemp()
    path = os.path.join(tmp_dir, '{}{}.{}'.format(name,
                                                  next(_counter),
                                                  ext))
    try:
        yield path
    finally:
        try:
            shutil.rmtree(tmp_dir)
        except OSError:
            if not ON_WIN:
                raise

@unittest.skipIf(condition=os.environ.get('CATE_DISABLE_PLOT_TESTS', None),
                 reason="skipped if CATE_DISABLE_PLOT_TESTS=1")
class TestPlotMap(TestCase):
    """
    Test plot_map() function
    """
    def test_plot_map(self):
        # Test the nominal functionality. This doesn't check that the plot is what is expected,
        # rather, it simply tests if it seems to have been created
        dataset = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.random.rand(5, 10, 2)),
            'second': (['lat', 'lon', 'time'], np.random.rand(5, 10, 2)),
            'lat': np.linspace(-89.5, 89.5, 5),
            'lon': np.linspace(-179.5, 179.5, 10)})

        with create_tmp_file('remove_me', 'png') as tmp_file:
            plot_map(dataset, file=tmp_file)
            self.assertTrue(os.path.isfile(tmp_file))

        # Test if an error is raised when an unsupported format is passed
        with create_tmp_file('remove_me', 'pgm') as tmp_file:
            with self.assertRaises(ValueError):
                plot_map(dataset, file=tmp_file)
            self.assertFalse(os.path.isfile(tmp_file))

        # Test if extents can be used
        with create_tmp_file('remove_me', 'pdf') as tmp_file:
            plot_map(dataset,
                     var='second',
                     region='-40.0, -20.0, 50.0, 60.0',
                     file=tmp_file)
            self.assertTrue(os.path.isfile(tmp_file))

    def test_plot_map_exceptions(self):
        # Test if the corner cases are detected without creating a plot for it.

        # Test value error is raised when passing an unexpected dataset type
        with create_tmp_file('remove_me', 'jpeg') as tmp_file:
            with self.assertRaises(NotImplementedError):
                plot_map([1, 2, 4], file=tmp_file)
            self.assertFalse(os.path.isfile(tmp_file))

        # Test the extensions bound checking
        dataset = xr.Dataset({
            'first': (['lat', 'lon', 'time'], np.random.rand(2, 4, 1)),
            'lat': np.linspace(-89.5, 89.5, 2),
            'lon': np.linspace(-179.5, 179.5, 4)})

        with self.assertRaises(ValueError):
            region = '-40.0, -95.0, 50.0, 60.0',
            plot_map(dataset, region=region)

        with self.assertRaises(ValueError):
            region = '-40.0, -20.0, 50.0, 95.0',
            plot_map(dataset, region=region)

        with self.assertRaises(ValueError):
            region = '-181.0, -20.0, 50.0, 60.0',
            plot_map(dataset, region=region)

        with self.assertRaises(ValueError):
            region = '-40.0, -20.0, 181.0, 60.0',
            plot_map(dataset, region=region)

        with self.assertRaises(ValueError):
            region = '-40.0, -20.0, 50.0, -25.0',
            plot_map(dataset, region=region)

        with self.assertRaises(ValueError):
            region = '-20.0, -20.0, -25.0, 60.0',
            plot_map(dataset, region=region)

    def test_registered(self):
        """
        Test nominal execution of the function as a registered operation.
        """
        pass

@unittest.skipIf(condition=os.environ.get('CATE_DISABLE_PLOT_TESTS', None),
                 reason="skipped if CATE_DISABLE_PLOT_TESTS=1")
class TestPlot(TestCase):
    """
    Test plot() function
    """
    def test_plot(self):
        # Test plot
        dataset = xr.Dataset({
            'first': np.random.rand(10)})

        with create_tmp_file('remove_me', 'jpg') as tmp_file:
            plot(dataset, 'first', file=tmp_file)
            self.assertTrue(os.path.isfile(tmp_file))

    def test_registered(self):
        """
        Test nominal execution of the function as a registered operation.
        """
        pass

@unittest.skipIf(condition=os.environ.get('CATE_DISABLE_PLOT_TESTS', None),
                 reason="skipped if CATE_DISABLE_PLOT_TESTS=1")
class TestPlotDataFrame(TestCase):
    """
    Test plot_dataframe() function.
    """
    def test_nominal(self):
        """
        Test nominal execution
        """
        pass

    def test_registered(self):
        """
        Test the method when run as a registered operation
        """
        pass
<|endoftext|>


<commit_before>#!/usr/bin/env python

from mpl_toolkits.mplot3d import Axes3D
import matplotlib.cm as cm
from sklearn.metrics import silhouette_score, silhouette_samples
import sklearn.cluster
import msmexplorer as msme
import mdtraj
import argparse
from matplotlib import pyplot as plt
import numpy as np
import seaborn as sns

parser = argparse.ArgumentParser(prog='cluster_ligands.py',
                                 formatter_class=argparse.RawDescriptionHelpFormatter,
                                 description='''version1''')

parser.add_argument("Trajectories", help="""An indefinite amount of AMBER
                    trajectories""", nargs="+")
parser.add_argument('-p', '--prmtop', type=str, required=True)
parser.add_argument('-st', '--stride', type=int, required=False, default=1)
parser.add_argument('-l', '--ligand_selection', type=str, required=False, default='all')
parser.add_argument('-o', '--out_file', type=str, required=False,
                    default='cluster_ligands')

def plot_com_matrix(com_matrix):
    f, axes = plt.subplots(3, 3, figsize=(15, 15))
    correspondance = {
        0: 'x',
        1: 'y',
        2: 'z'
    }
    # axes is a np.array of shape (3, 3)
    for i in range(3):
        for j in range(3):
            ax = axes[i][j]
            if i >= j:
                # we'll plot only the lower half
                if i == j:
                    sns.kdeplot(com_matrix[:, i], ax=ax)
                    ax.set(xlabel='%s (Å)' % correspondance[i],
                           ylabel='Density')
                else:
                    ax.set(xlabel='%s (Å)' % correspondance[j],
                           ylabel='%s (Å)' % correspondance[i])
                    msme.plot_free_energy(
                        com_matrix,
                        obs=(j, i),
                        ax=ax,
                        shade=True,
                        n_levels=5,
                        vmin=-1e-12,
                        cmap='viridis'
                    )
            else:
                ax.set_axis_off()
    return f

def unique_rows(a):
    a = np.ascontiguousarray(a)
    unique_a = np.unique(a.view([('', a.dtype)] * a.shape[1]))
    return unique_a.view(a.dtype).reshape((unique_a.shape[0], a.shape[1]))

def report_clusters(X):
    fig_list = []
    avg_sils = []
    for n_clusters in [2, 3, 4, 5, 6]:
        fig = plt.figure()
        ax1 = plt.subplot(121)
        ax2 = plt.subplot(122, projection='3d')
        ax2.w_xaxis.set_pane_color((1.0, 1.0, 1.0, 1.0))
        ax2.w_yaxis.set_pane_color((1.0, 1.0, 1.0, 1.0))
        ax2.w_zaxis.set_pane_color((1.0, 1.0, 1.0, 1.0))
        fig.set_size_inches(18, 7)
        ax1.set_xlim([-0.2, 1])
        ax1.set_ylim([0, len(X) + (n_clusters + 1) * 10])
        clusterer = sklearn.cluster.KMeans(n_clusters=n_clusters)
        cluster_labels = clusterer.fit_predict(X)
        silhouette_avg = silhouette_score(X, cluster_labels)
        avg_sils.append(silhouette_avg)
        sample_silhouette_values = silhouette_samples(X, cluster_labels)
        y_lower = 10
        for i in range(n_clusters):
            # Aggregate the silhouette scores for samples belonging to
            # cluster i, and sort them
            ith_cluster_silhouette_values = \
                sample_silhouette_values[cluster_labels == i]

            ith_cluster_silhouette_values.sort()

            size_cluster_i = ith_cluster_silhouette_values.shape[0]
            y_upper = y_lower + size_cluster_i

            color = cm.spectral(float(i) / n_clusters)
            ax1.fill_betweenx(np.arange(y_lower, y_upper),
                              0, ith_cluster_silhouette_values,
                              facecolor=color, edgecolor=color, alpha=0.7)

            # Label the silhouette plots with their cluster numbers at the middle
            ax1.text(-0.05, y_lower + 0.5 * size_cluster_i, str(i))

            # Compute the new y_lower for next plot
            y_lower = y_upper + 10  # 10 for the 0 samples

        ax1.set_title("Silhouette plot")
        ax1.set_xlabel("Silhouette coefficient")
        ax1.set_ylabel("Cluster label")

        # The vertical line for average silhouette score of all the values
        ax1.axvline(x=silhouette_avg, color="red", linestyle="--")
        ax1.annotate('Average silhouette: %.2f' % silhouette_avg, xy=(0.80, 0.95), xycoords='axes fraction')
        ax1.set_yticks([])  # Clear the yaxis labels / ticks
        ax1.set_xticks([-0.2, -0.1, 0, 0.2, 0.4, 0.6, 0.8, 1])

        # 2nd Plot showing the actual clusters formed
        colors = cm.spectral(cluster_labels.astype(float) / n_clusters)
        ax2.scatter(X[:, 0], X[:, 1], X[:, 2], marker='.', s=30, lw=0, alpha=0.1,
                    c=colors)

        # Labeling the clusters
        centers = clusterer.cluster_centers_

        unique_colors = unique_rows(colors)
        for i, c in enumerate(zip(centers, unique_colors)):
            c, col = c[0], c[1]
            # Draw white circles at cluster centers
            ax2.scatter(c[0], c[1], c[2], marker='o', alpha=1, s=150, color='white')
            ax2.scatter(c[0], c[1], c[2], marker='$%d$' % i, alpha=1, s=100, color=col)

        ax2.set_xlabel("x (nm)")
        ax2.set_ylabel("y (nm)")
        ax2.set_zlabel("y (nm)")

        fig_list.append((fig, n_clusters))
    return fig_list

if __name__ == '__main__':
    args = parser.parse_args()
    print(args)
    top = mdtraj.load_prmtop(args.prmtop)
    traj = mdtraj.load([t for t in args.Trajectories], top=args.prmtop,
                       stride=args.stride)
    # Center all coordinates, makes center of geometry of the system (0, 0, 0)
    # traj.center_coordinates()
    # Superpose trajectory onto first frame
    # This is exactly like the 'hold selection steady' command in Chimera
    traj.superpose(traj, 0)
    # Create separate trajectory for the ligand
    ligand_indices = top.select(args.ligand_selection)
    lig_traj = traj.atom_slice(ligand_indices)
    center_mass_ligand = mdtraj.compute_center_of_mass(lig_traj)
    time = np.linspace(0, traj.timestep * traj.n_frames, num=traj.n_frames)
    f = plot_com_matrix(center_mass_ligand)
    f.savefig(args.out_file + '.pdf')
    fig_list = report_clusters(center_mass_ligand)
    for fig in fig_list:
        f, n_clusters = fig[0], fig[1]
        f.savefig(args.out_file + str(n_clusters) + '.pdf')
<commit_msg>z label fix<commit_after>#!/usr/bin/env python

from mpl_toolkits.mplot3d import Axes3D
import matplotlib.cm as cm
from sklearn.metrics import silhouette_score, silhouette_samples
import sklearn.cluster
import msmexplorer as msme
import mdtraj
import argparse
from matplotlib import pyplot as plt
import numpy as np
import seaborn as sns

parser = argparse.ArgumentParser(prog='cluster_ligands.py',
                                 formatter_class=argparse.RawDescriptionHelpFormatter,
                                 description='''version1''')

parser.add_argument("Trajectories", help="""An indefinite amount of AMBER
                    trajectories""", nargs="+")
parser.add_argument('-p', '--prmtop', type=str, required=True)
parser.add_argument('-st', '--stride', type=int, required=False, default=1)
parser.add_argument('-l', '--ligand_selection', type=str, required=False, default='all')
parser.add_argument('-o', '--out_file', type=str, required=False,
                    default='cluster_ligands')

def plot_com_matrix(com_matrix):
    f, axes = plt.subplots(3, 3, figsize=(15, 15))
    correspondance = {
        0: 'x',
        1: 'y',
        2: 'z'
    }
    # axes is a np.array of shape (3, 3)
    for i in range(3):
        for j in range(3):
            ax = axes[i][j]
            if i >= j:
                # we'll plot only the lower half
                if i == j:
                    sns.kdeplot(com_matrix[:, i], ax=ax)
                    ax.set(xlabel='%s (Å)' % correspondance[i],
                           ylabel='Density')
                else:
                    ax.set(xlabel='%s (Å)' % correspondance[j],
                           ylabel='%s (Å)' % correspondance[i])
                    msme.plot_free_energy(
                        com_matrix,
                        obs=(j, i),
                        ax=ax,
                        shade=True,
                        n_levels=5,
                        vmin=-1e-12,
                        cmap='viridis'
                    )
            else:
                ax.set_axis_off()
    return f

def unique_rows(a):
    a = np.ascontiguousarray(a)
    unique_a = np.unique(a.view([('', a.dtype)] * a.shape[1]))
    return unique_a.view(a.dtype).reshape((unique_a.shape[0], a.shape[1]))

def report_clusters(X):
    fig_list = []
    avg_sils = []
    for n_clusters in [2, 3, 4, 5, 6]:
        fig = plt.figure()
        ax1 = plt.subplot(121)
        ax2 = plt.subplot(122, projection='3d')
        ax2.w_xaxis.set_pane_color((1.0, 1.0, 1.0, 1.0))
        ax2.w_yaxis.set_pane_color((1.0, 1.0, 1.0, 1.0))
        ax2.w_zaxis.set_pane_color((1.0, 1.0, 1.0, 1.0))
        fig.set_size_inches(18, 7)
        ax1.set_xlim([-0.2, 1])
        ax1.set_ylim([0, len(X) + (n_clusters + 1) * 10])
        clusterer = sklearn.cluster.KMeans(n_clusters=n_clusters)
        cluster_labels = clusterer.fit_predict(X)
        silhouette_avg = silhouette_score(X, cluster_labels)
        avg_sils.append(silhouette_avg)
        sample_silhouette_values = silhouette_samples(X, cluster_labels)
        y_lower = 10
        for i in range(n_clusters):
            # Aggregate the silhouette scores for samples belonging to
            # cluster i, and sort them
            ith_cluster_silhouette_values = \
                sample_silhouette_values[cluster_labels == i]

            ith_cluster_silhouette_values.sort()

            size_cluster_i = ith_cluster_silhouette_values.shape[0]
            y_upper = y_lower + size_cluster_i

            color = cm.spectral(float(i) / n_clusters)
            ax1.fill_betweenx(np.arange(y_lower, y_upper),
                              0, ith_cluster_silhouette_values,
                              facecolor=color, edgecolor=color, alpha=0.7)

            # Label the silhouette plots with their cluster numbers at the middle
            ax1.text(-0.05, y_lower + 0.5 * size_cluster_i, str(i))

            # Compute the new y_lower for next plot
            y_lower = y_upper + 10  # 10 for the 0 samples

        ax1.set_title("Silhouette plot")
        ax1.set_xlabel("Silhouette coefficient")
        ax1.set_ylabel("Cluster label")

        # The vertical line for average silhouette score of all the values
        ax1.axvline(x=silhouette_avg, color="red", linestyle="--")
        ax1.annotate('Average silhouette: %.2f' % silhouette_avg, xy=(0.80, 0.95), xycoords='axes fraction')
        ax1.set_yticks([])  # Clear the yaxis labels / ticks
        ax1.set_xticks([-0.2, -0.1, 0, 0.2, 0.4, 0.6, 0.8, 1])

        # 2nd Plot showing the actual clusters formed
        colors = cm.spectral(cluster_labels.astype(float) / n_clusters)
        ax2.scatter(X[:, 0], X[:, 1], X[:, 2], marker='.', s=30, lw=0, alpha=0.1,
                    c=colors)

        # Labeling the clusters
        centers = clusterer.cluster_centers_

        unique_colors = unique_rows(colors)
        for i, c in enumerate(zip(centers, unique_colors)):
            c, col = c[0], c[1]
            # Draw white circles at cluster centers
            ax2.scatter(c[0], c[1], c[2], marker='o', alpha=1, s=150, color='white')
            ax2.scatter(c[0], c[1], c[2], marker='$%d$' % i, alpha=1, s=100, color=col)

        ax2.set_xlabel("x (nm)")
        ax2.set_ylabel("y (nm)")
        ax2.set_zlabel("z (nm)")

        fig_list.append((fig, n_clusters))
    return fig_list

if __name__ == '__main__':
    args = parser.parse_args()
    print(args)
    top = mdtraj.load_prmtop(args.prmtop)
    traj = mdtraj.load([t for t in args.Trajectories], top=args.prmtop,
                       stride=args.stride)
    # Center all coordinates, makes center of geometry of the system (0, 0, 0)
    # traj.center_coordinates()
    # Superpose trajectory onto first frame
    # This is exactly like the 'hold selection steady' command in Chimera
    traj.superpose(traj, 0)
    # Create separate trajectory for the ligand
    ligand_indices = top.select(args.ligand_selection)
    lig_traj = traj.atom_slice(ligand_indices)
    center_mass_ligand = mdtraj.compute_center_of_mass(lig_traj)
    time = np.linspace(0, traj.timestep * traj.n_frames, num=traj.n_frames)
    f = plot_com_matrix(center_mass_ligand)
    f.savefig(args.out_file + '.pdf')
    fig_list = report_clusters(center_mass_ligand)
    for fig in fig_list:
        f, n_clusters = fig[0], fig[1]
        f.savefig(args.out_file + str(n_clusters) + '.pdf')
<|endoftext|>


<commit_before># Copyright 2014 Bevbot LLC, All Rights Reserved
#
# This file is part of the Pykeg package of the Kegbot project.
# For more information on Pykeg or Kegbot, see http://kegbot.org/
#
# Pykeg is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 2 of the License, or
# (at your option) any later version.
#
# Pykeg is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Pykeg.  If not, see <http://www.gnu.org/licenses/>.

"""Kegweb main views."""

from django.contrib import messages
from django.shortcuts import get_object_or_404
from django.shortcuts import render_to_response
from django.shortcuts import redirect
from django.template import RequestContext
from django.views.decorators.cache import cache_page
from django.views.generic.dates import ArchiveIndexView
from django.views.generic.dates import DateDetailView
from django.views.generic.dates import DayArchiveView
from django.views.generic.dates import MonthArchiveView
from django.views.generic.dates import YearArchiveView
from django.views.generic.list import ListView
from django.core.paginator import Paginator, EmptyPage, PageNotAnInteger

from pykeg.core import models
from pykeg.web.kegweb import forms

### main views

@cache_page(30)
def index(request):
    context = RequestContext(request)

    context['taps'] = models.KegTap.objects.all()
    context['events'] = models.SystemEvent.objects.timeline()[:20]
    sessions = models.DrinkingSession.objects.all().order_by('-id')[:10]
    context['sessions'] = sessions

    if sessions:
        last_session = sessions[0]
        context['most_recent_session'] = last_session
        if sessions and last_session.IsActiveNow():
            context['current_session'] = last_session

    return render_to_response('index.html', context_instance=context)

@cache_page(30)
def system_stats(request):
    stats = models.KegbotSite.get().get_stats()
    context = RequestContext(request, {
        'stats': stats,
    })

    top_drinkers = []
    for username, vol in stats.get('volume_by_drinker', {}).iteritems():
        try:
            user = models.User.objects.get(username=username)
        except models.User.DoesNotExist:
            continue  # should not happen
        top_drinkers.append((vol, user))
    top_drinkers.sort(reverse=True)

    largest_session_id = stats.get('largest_session', {}).get('session_id', None)
    if largest_session_id:
        try:
            context['largest_session'] = models.DrinkingSession.objects.get(pk=largest_session_id)
        except models.DrinkingSession.DoesNotExist:
            # Stats out of date.
            pass

    context['top_drinkers'] = top_drinkers[:10]

    return render_to_response('kegweb/system-stats.html', context_instance=context)

### object lists and detail (generic views)

def user_detail(request, username):
    user = get_object_or_404(models.User, username=username, is_active=True)
    stats = user.get_stats()
    drinks = user.drinks.all()

    context = RequestContext(request, {
        'drinks': drinks,
        'stats': stats,
        'drinker': user})

    largest_session_id = stats.get('largest_session', {}).get('session_id', None)
    if largest_session_id:
        context['largest_session'] = models.DrinkingSession.objects.get(pk=largest_session_id)

    return render_to_response('kegweb/drinker_detail.html', context_instance=context)

class KegListView(ListView):
    model = models.Keg
    template_name = 'kegweb/keg_list.html'
    context_object_name = 'kegs'
    paginate_by = 10

    def get_queryset(self):
        return models.Keg.objects.all().order_by('-id')

def fullscreen(request):
    context = RequestContext(request)
    taps = models.KegTap.objects.all()
    active_taps = [t for t in taps if t.current_keg]
    pages = [active_taps[i:i + 4] for i in range(0, len(active_taps), 4)]
    context['pages'] = pages

    return render_to_response('kegweb/fullscreen.html', context_instance=context)

@cache_page(30)
def keg_detail(request, keg_id):
    keg = get_object_or_404(models.Keg, id=keg_id)
    sessions = keg.get_sessions()
    last_session = sessions[:1]

    context = RequestContext(request, {
        'keg': keg,
        'stats': keg.get_stats(),
        'sessions': sessions,
        'last_session': last_session})
    return render_to_response('kegweb/keg_detail.html', context_instance=context)

def short_drink_detail(request, drink_id):
    return redirect('kb-drink', drink_id=str(drink_id), permanent=True)

def short_session_detail(request, session_id):
    session = get_object_or_404(models.DrinkingSession, id=session_id)
    url = session.get_absolute_url()
    return redirect(url, permanent=True)

def drink_detail(request, drink_id):
    drink = get_object_or_404(models.Drink, id=drink_id)
    context = RequestContext(request, {'drink': drink})

    can_delete = (request.user == drink.user) or request.user.is_staff

    if can_delete:
        picture_form = forms.DeletePictureForm(initial={'picture': drink.picture})
    else:
        picture_form = None

    if request.method == 'POST':
        if can_delete:
            picture_form = forms.DeletePictureForm(request.POST)
            if picture_form.is_valid():
                drink.picture.erase_and_delete()
                picture_form = None
                messages.success(request, 'Erased image.')
            else:
                messages.error(request, 'request not valid: ' + str(picture_form.errors))
        else:
            messages.error(request, 'No permission to delete picture.')
        return redirect('kb-drink', drink_id=str(drink_id))

    context['picture_form'] = picture_form
    return render_to_response('kegweb/drink_detail.html', context_instance=context)

def drinker_sessions(request, username):
    user = get_object_or_404(models.User, username=username, is_active=True)
    stats = user.get_stats()
    drinks = user.drinks.all()

    chunks = models.Stats.objects.filter(
        user=user, keg__isnull=True, session__isnull=False, is_first=True
    ).order_by('-id').select_related('session')

    paginator = Paginator(chunks, 5)

    page = request.GET.get('page')
    try:
        chunks = paginator.page(page)
    except PageNotAnInteger:
        chunks = paginator.page(1)
    except EmptyPage:
        chunks = paginator.page(paginator.num_pages)

    context = RequestContext(request, {
        'drinks': drinks,
        'chunks': chunks,
        'stats': stats,
        'drinker': user})

    return render_to_response('kegweb/drinker_sessions.html', context_instance=context)

def keg_sessions(request, keg_id):
    keg = get_object_or_404(models.Keg, id=keg_id)
    sessions = keg.get_sessions()

    paginator = Paginator(sessions, 5)

    page = request.GET.get('page')
    try:
        sessions = paginator.page(page)
    except PageNotAnInteger:
        sessions = paginator.page(1)
    except EmptyPage:
        sessions = paginator.page(paginator.num_pages)

    context = RequestContext(request, {
        'keg': keg,
        'stats': keg.get_stats(),
        'sessions': sessions})
    return render_to_response('kegweb/keg_sessions.html', context_instance=context)

class SessionArchiveIndexView(ArchiveIndexView):
    model = models.DrinkingSession
    date_field = 'start_time'
    template_name = 'kegweb/drinkingsession_archive.html'
    context_object_name = 'sessions'
    paginate_by = 20

class SessionYearArchiveView(YearArchiveView):
    model = models.DrinkingSession
    date_field = 'start_time'
    template_name = 'kegweb/drinkingsession_archive_year.html'
    make_object_list = True
    context_object_name = 'sessions'
    paginate_by = 20

class SessionMonthArchiveView(MonthArchiveView):
    model = models.DrinkingSession
    date_field = 'start_time'
    template_name = 'kegweb/drinkingsession_archive_month.html'
    make_object_list = True
    context_object_name = 'sessions'
    paginate_by = 20

class SessionDayArchiveView(DayArchiveView):
    model = models.DrinkingSession
    date_field = 'start_time'
    template_name = 'kegweb/drinkingsession_archive_day.html'
    make_object_list = True
    context_object_name = 'sessions'
    paginate_by = 20

class SessionDateDetailView(DateDetailView):
    model = models.DrinkingSession
    date_field = 'start_time'
    template_name = 'kegweb/session_detail.html'
    context_object_name = 'session'

    def get_context_data(self, **kwargs):
        """Adds `stats` to the context."""
        ret = super(SessionDateDetailView, self).get_context_data(**kwargs)
        stats = ret[self.context_object_name].get_stats()
        ret['stats'] = stats
        ret['kegs'] = [models.Keg.objects.get(pk=pk) for pk in stats.get('keg_ids', [])]
        return ret
<commit_msg>Don't require user to be active to view details and sessions.<commit_after># Copyright 2014 Bevbot LLC, All Rights Reserved
#
# This file is part of the Pykeg package of the Kegbot project.
# For more information on Pykeg or Kegbot, see http://kegbot.org/
#
# Pykeg is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 2 of the License, or
# (at your option) any later version.
#
# Pykeg is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Pykeg.  If not, see <http://www.gnu.org/licenses/>.

"""Kegweb main views."""

from django.contrib import messages
from django.shortcuts import get_object_or_404
from django.shortcuts import render_to_response
from django.shortcuts import redirect
from django.template import RequestContext
from django.views.decorators.cache import cache_page
from django.views.generic.dates import ArchiveIndexView
from django.views.generic.dates import DateDetailView
from django.views.generic.dates import DayArchiveView
from django.views.generic.dates import MonthArchiveView
from django.views.generic.dates import YearArchiveView
from django.views.generic.list import ListView
from django.core.paginator import Paginator, EmptyPage, PageNotAnInteger

from pykeg.core import models
from pykeg.web.kegweb import forms

### main views

@cache_page(30)
def index(request):
    context = RequestContext(request)

    context['taps'] = models.KegTap.objects.all()
    context['events'] = models.SystemEvent.objects.timeline()[:20]
    sessions = models.DrinkingSession.objects.all().order_by('-id')[:10]
    context['sessions'] = sessions

    if sessions:
        last_session = sessions[0]
        context['most_recent_session'] = last_session
        if sessions and last_session.IsActiveNow():
            context['current_session'] = last_session

    return render_to_response('index.html', context_instance=context)

@cache_page(30)
def system_stats(request):
    stats = models.KegbotSite.get().get_stats()
    context = RequestContext(request, {
        'stats': stats,
    })

    top_drinkers = []
    for username, vol in stats.get('volume_by_drinker', {}).iteritems():
        try:
            user = models.User.objects.get(username=username)
        except models.User.DoesNotExist:
            continue  # should not happen
        top_drinkers.append((vol, user))
    top_drinkers.sort(reverse=True)

    largest_session_id = stats.get('largest_session', {}).get('session_id', None)
    if largest_session_id:
        try:
            context['largest_session'] = models.DrinkingSession.objects.get(pk=largest_session_id)
        except models.DrinkingSession.DoesNotExist:
            # Stats out of date.
            pass

    context['top_drinkers'] = top_drinkers[:10]

    return render_to_response('kegweb/system-stats.html', context_instance=context)

### object lists and detail (generic views)

def user_detail(request, username):
    user = get_object_or_404(models.User, username=username)
    stats = user.get_stats()
    drinks = user.drinks.all()

    context = RequestContext(request, {
        'drinks': drinks,
        'stats': stats,
        'drinker': user})

    largest_session_id = stats.get('largest_session', {}).get('session_id', None)
    if largest_session_id:
        context['largest_session'] = models.DrinkingSession.objects.get(pk=largest_session_id)

    return render_to_response('kegweb/drinker_detail.html', context_instance=context)

class KegListView(ListView):
    model = models.Keg
    template_name = 'kegweb/keg_list.html'
    context_object_name = 'kegs'
    paginate_by = 10

    def get_queryset(self):
        return models.Keg.objects.all().order_by('-id')

def fullscreen(request):
    context = RequestContext(request)
    taps = models.KegTap.objects.all()
    active_taps = [t for t in taps if t.current_keg]
    pages = [active_taps[i:i + 4] for i in range(0, len(active_taps), 4)]
    context['pages'] = pages

    return render_to_response('kegweb/fullscreen.html', context_instance=context)

@cache_page(30)
def keg_detail(request, keg_id):
    keg = get_object_or_404(models.Keg, id=keg_id)
    sessions = keg.get_sessions()
    last_session = sessions[:1]

    context = RequestContext(request, {
        'keg': keg,
        'stats': keg.get_stats(),
        'sessions': sessions,
        'last_session': last_session})
    return render_to_response('kegweb/keg_detail.html', context_instance=context)

def short_drink_detail(request, drink_id):
    return redirect('kb-drink', drink_id=str(drink_id), permanent=True)

def short_session_detail(request, session_id):
    session = get_object_or_404(models.DrinkingSession, id=session_id)
    url = session.get_absolute_url()
    return redirect(url, permanent=True)

def drink_detail(request, drink_id):
    drink = get_object_or_404(models.Drink, id=drink_id)
    context = RequestContext(request, {'drink': drink})

    can_delete = (request.user == drink.user) or request.user.is_staff

    if can_delete:
        picture_form = forms.DeletePictureForm(initial={'picture': drink.picture})
    else:
        picture_form = None

    if request.method == 'POST':
        if can_delete:
            picture_form = forms.DeletePictureForm(request.POST)
            if picture_form.is_valid():
                drink.picture.erase_and_delete()
                picture_form = None
                messages.success(request, 'Erased image.')
            else:
                messages.error(request, 'request not valid: ' + str(picture_form.errors))
        else:
            messages.error(request, 'No permission to delete picture.')
        return redirect('kb-drink', drink_id=str(drink_id))

    context['picture_form'] = picture_form
    return render_to_response('kegweb/drink_detail.html', context_instance=context)

def drinker_sessions(request, username):
    user = get_object_or_404(models.User, username=username)
    stats = user.get_stats()
    drinks = user.drinks.all()

    chunks = models.Stats.objects.filter(
        user=user, keg__isnull=True, session__isnull=False, is_first=True
    ).order_by('-id').select_related('session')

    paginator = Paginator(chunks, 5)

    page = request.GET.get('page')
    try:
        chunks = paginator.page(page)
    except PageNotAnInteger:
        chunks = paginator.page(1)
    except EmptyPage:
        chunks = paginator.page(paginator.num_pages)

    context = RequestContext(request, {
        'drinks': drinks,
        'chunks': chunks,
        'stats': stats,
        'drinker': user})

    return render_to_response('kegweb/drinker_sessions.html', context_instance=context)

def keg_sessions(request, keg_id):
    keg = get_object_or_404(models.Keg, id=keg_id)
    sessions = keg.get_sessions()

    paginator = Paginator(sessions, 5)

    page = request.GET.get('page')
    try:
        sessions = paginator.page(page)
    except PageNotAnInteger:
        sessions = paginator.page(1)
    except EmptyPage:
        sessions = paginator.page(paginator.num_pages)

    context = RequestContext(request, {
        'keg': keg,
        'stats': keg.get_stats(),
        'sessions': sessions})
    return render_to_response('kegweb/keg_sessions.html', context_instance=context)

class SessionArchiveIndexView(ArchiveIndexView):
    model = models.DrinkingSession
    date_field = 'start_time'
    template_name = 'kegweb/drinkingsession_archive.html'
    context_object_name = 'sessions'
    paginate_by = 20

class SessionYearArchiveView(YearArchiveView):
    model = models.DrinkingSession
    date_field = 'start_time'
    template_name = 'kegweb/drinkingsession_archive_year.html'
    make_object_list = True
    context_object_name = 'sessions'
    paginate_by = 20

class SessionMonthArchiveView(MonthArchiveView):
    model = models.DrinkingSession
    date_field = 'start_time'
    template_name = 'kegweb/drinkingsession_archive_month.html'
    make_object_list = True
    context_object_name = 'sessions'
    paginate_by = 20

class SessionDayArchiveView(DayArchiveView):
    model = models.DrinkingSession
    date_field = 'start_time'
    template_name = 'kegweb/drinkingsession_archive_day.html'
    make_object_list = True
    context_object_name = 'sessions'
    paginate_by = 20

class SessionDateDetailView(DateDetailView):
    model = models.DrinkingSession
    date_field = 'start_time'
    template_name = 'kegweb/session_detail.html'
    context_object_name = 'session'

    def get_context_data(self, **kwargs):
        """Adds `stats` to the context."""
        ret = super(SessionDateDetailView, self).get_context_data(**kwargs)
        stats = ret[self.context_object_name].get_stats()
        ret['stats'] = stats
        ret['kegs'] = [models.Keg.objects.get(pk=pk) for pk in stats.get('keg_ids', [])]
        return ret
<|endoftext|>


<commit_before>import unittest
from PandasRake.rake import Rake
from PandasRake.datasets import load_people

class TestPandasRake(unittest.TestCase):

    def setUp(self):
        self.recodes = {"Age": lambda age: age / 10,
                        "MaritalStatus": lambda m: m / 2}
        self.target_pops = {0: .5, 1: .2}
        self.df = load_people()

    def test_smoke(self):
        r = Rake(self.df, self.recodes, self.target_pops)
        self.assertIsNotNone(r)

    def test_recode(self):
        r = Rake(self.df, self.recodes, self.target_pops)
        r.recode()
        self.assertEqual(r.df['Age'].values[0], 8)
        self.assertEqual(r.df['MaritalStatus'].values[2], 2)

    def test_rake(self):
        pass
        #r = Rake(self.df, self.recodes, self.target_pops)
        # r.recode()
        #wt = r.rake()
        #self.assertEqual(len(r.weights), len(wt))
<commit_msg>fixed comment block<commit_after>import unittest
from PandasRake.rake import Rake
from PandasRake.datasets import load_people

class TestPandasRake(unittest.TestCase):

    def setUp(self):
        self.recodes = {"Age": lambda age: age / 10,
                        "MaritalStatus": lambda m: m / 2}
        self.target_pops = {0: .5, 1: .2}
        self.df = load_people()

    def test_smoke(self):
        r = Rake(self.df, self.recodes, self.target_pops)
        self.assertIsNotNone(r)

    def test_recode(self):
        r = Rake(self.df, self.recodes, self.target_pops)
        r.recode()
        self.assertEqual(r.df['Age'].values[0], 8)
        self.assertEqual(r.df['MaritalStatus'].values[2], 2)

    def test_rake(self):
        pass
        # r = Rake(self.df, self.recodes, self.target_pops)
        # r.recode()
        # wt = r.rake()
        # self.assertEqual(len(r.weights), len(wt))
<|endoftext|>


<commit_before>"""
 Python wrapper for libui.

"""

import ctypes
from . import clibui

class uiCheckbox(ctypes.Structure):
    """Wrapper for the uiCheckbox C struct."""

    pass

def uiCheckboxPointer(obj):
    """
    Casts an object to uiCheckbox pointer type.

    :param obj: a generic object
    :return: uiCheckbox
    """

    return ctypes.cast(obj, ctypes.POINTER(uiCheckbox))

# - char *uiCheckboxText(uiCheckbox *c);
def uiCheckboxText(checkbox):
    """
    Returns the text of the checkbox.

    :param checkbox: uiCheckbox
    :return: string
    """

    clibui.uiCheckboxText.restype = ctypes.c_char_p
    text = clibui.uiCheckboxText(checkbox)

    return text.decode()

# - void uiCheckboxSetText(uiCheckbox *c, const char *text);
def uiCheckboxSetText(checkbox, text):
    """
    Sets the text of the checkbox.

    :param checkbox: uiCheckbox
    :param text: string
    :return: None
    """

    clibui.uiCheckboxSetText(checkbox, bytes(text, 'utf-8'))

# - void uiCheckboxOnToggled(uiCheckbox *c, void (*f)(uiCheckbox *c, void *data), void *data);
def uiCheckboxOnToggled(checkbox, callback, data):
    """
    Executes the callback function on checkbox toggle.

    :param checkbox: uiCheckbox
    :param callback: function
    :param data: data
    :return: reference to C callback function
    """
    c_type = ctypes.CFUNCTYPE(
        ctypes.c_int, ctypes.POINTER(uiCheckbox), ctypes.c_void_p)
    c_callback = c_type(callback)

    clibui.uiButtonOnClicked(checkbox, c_callback, data)

    return c_callback

# - int uiCheckboxChecked(uiCheckbox *c);
def uiCheckboxChecked(checkbox):
    """
    Returns whether the checkbox is checked or not.

    :param checkbox: uiCheckbox
    :return: string
    """

    return bool(clibui.uiCheckboxChecked(checkbox))

# - void uiCheckboxSetChecked(uiCheckbox *c, int checked);
def uiCheckboxSetChecked(checkbox, checked):
    """
    Sets whether the checkbox is checked or not.

    :param checkbox: uiCheckbox
    :param checked: bool
    :return: None
    """

    clibui.uiCheckboxSetChecked(checkbox, checked)

# - uiCheckbox *uiNewCheckbox(const char *text);
def uiNewCheckbox(text):
    """
    Creates a new checkbox.

    :return: uiCheckbox
    """

    # Set return type
    clibui.uiNewCheckbox.restype = ctypes.POINTER(uiCheckbox)

    return clibui.uiNewCheckbox(bytes(text, 'utf-8'))
<commit_msg>fixed wrong function call<commit_after>"""
 Python wrapper for libui.

"""

import ctypes
from . import clibui

class uiCheckbox(ctypes.Structure):
    """Wrapper for the uiCheckbox C struct."""

    pass

def uiCheckboxPointer(obj):
    """
    Casts an object to uiCheckbox pointer type.

    :param obj: a generic object
    :return: uiCheckbox
    """

    return ctypes.cast(obj, ctypes.POINTER(uiCheckbox))

# - char *uiCheckboxText(uiCheckbox *c);
def uiCheckboxText(checkbox):
    """
    Returns the text of the checkbox.

    :param checkbox: uiCheckbox
    :return: string
    """

    clibui.uiCheckboxText.restype = ctypes.c_char_p
    text = clibui.uiCheckboxText(checkbox)

    return text.decode()

# - void uiCheckboxSetText(uiCheckbox *c, const char *text);
def uiCheckboxSetText(checkbox, text):
    """
    Sets the text of the checkbox.

    :param checkbox: uiCheckbox
    :param text: string
    :return: None
    """

    clibui.uiCheckboxSetText(checkbox, bytes(text, 'utf-8'))

# - void uiCheckboxOnToggled(uiCheckbox *c, void (*f)(uiCheckbox *c, void *data), void *data);
def uiCheckboxOnToggled(checkbox, callback, data):
    """
    Executes the callback function on checkbox toggle.

    :param checkbox: uiCheckbox
    :param callback: function
    :param data: data
    :return: reference to C callback function
    """
    c_type = ctypes.CFUNCTYPE(
        ctypes.c_int, ctypes.POINTER(uiCheckbox), ctypes.c_void_p)
    c_callback = c_type(callback)

    clibui.uiCheckboxOnToggled(checkbox, c_callback, data)

    return c_callback

# - int uiCheckboxChecked(uiCheckbox *c);
def uiCheckboxChecked(checkbox):
    """
    Returns whether the checkbox is checked or not.

    :param checkbox: uiCheckbox
    :return: string
    """

    return bool(clibui.uiCheckboxChecked(checkbox))

# - void uiCheckboxSetChecked(uiCheckbox *c, int checked);
def uiCheckboxSetChecked(checkbox, checked):
    """
    Sets whether the checkbox is checked or not.

    :param checkbox: uiCheckbox
    :param checked: bool
    :return: None
    """

    clibui.uiCheckboxSetChecked(checkbox, checked)

# - uiCheckbox *uiNewCheckbox(const char *text);
def uiNewCheckbox(text):
    """
    Creates a new checkbox.

    :return: uiCheckbox
    """

    # Set return type
    clibui.uiNewCheckbox.restype = ctypes.POINTER(uiCheckbox)

    return clibui.uiNewCheckbox(bytes(text, 'utf-8'))
<|endoftext|>


<commit_before>import unittest

import mock
from twisted.internet import protocol
from twisted.test import proto_helpers

from txrudp import heap, packet, rudp

class TestConnectionManagerAPI(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        cls.public_ip = '123.45.67.89'
        cls.port = 12345
        cls.addr = (cls.public_ip, cls.port)

    def _make_cm_with_mocks(self):
        cf = mock.Mock()
        return rudp.ConnectionMultiplexer(cf, self.public_ip)

    def test_default_init(self):
        cf = mock.Mock()
        cm = rudp.ConnectionMultiplexer(cf, self.public_ip)
        self.assertIsInstance(cm, protocol.DatagramProtocol)
        self.assertEqual(cm.public_ip, self.public_ip)
        self.assertFalse(cm.relaying)
        self.assertEqual(len(cm), 0)

    def test_full_init(self):
        cf = mock.Mock()
        cm = rudp.ConnectionMultiplexer(
            connection_factory=cf,
            public_ip='192.168.1.1',
            relaying=True
        )
        self.assertIsInstance(cm, protocol.DatagramProtocol)
        self.assertEqual(cm.public_ip, '192.168.1.1')
        self.assertTrue(cm.relaying)
        self.assertIsNone(cm.transport)
        self.assertEqual(len(cm), 0)

    def test_make_connection(self):
        transport = proto_helpers.StringTransportWithDisconnection()
        cm = self._make_cm_with_mocks()
        cm.makeConnection(transport)
        self.assertIs(cm.transport, transport)

    def test_get_nonexistent_connection(self):
        cm = self._make_cm_with_mocks()
        self.assertNotIn(self.addr, cm)
        with self.assertRaises(KeyError):
            con = cm[self.addr]

    def test_set_and_get_new_connection(self):
        cm = self._make_cm_with_mocks()
        mock_connection = mock.Mock()
        cm[self.addr] = mock_connection
        self.assertIn(self.addr, cm)
        self.assertEqual(cm[self.addr], mock_connection)

    def test_set_existent_connection(self):
        pass
<commit_msg>Stricten mocks used in tests.<commit_after>import unittest

import mock
from twisted.internet import protocol
from twisted.test import proto_helpers

from txrudp import connection, heap, packet, rudp

class TestConnectionManagerAPI(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        cls.public_ip = '123.45.67.89'
        cls.port = 12345
        cls.addr = (cls.public_ip, cls.port)

    def _make_cm_with_mocks(self):
        cf = mock.Mock(spec_set=connection.RUDPConnectionFactory)
        return rudp.ConnectionMultiplexer(cf, self.public_ip)

    def test_default_init(self):
        cf = mock.Mock()
        cm = rudp.ConnectionMultiplexer(cf, self.public_ip)
        self.assertIsInstance(cm, protocol.DatagramProtocol)
        self.assertEqual(cm.public_ip, self.public_ip)
        self.assertFalse(cm.relaying)
        self.assertEqual(len(cm), 0)

    def test_full_init(self):
        cf = mock.Mock()
        cm = rudp.ConnectionMultiplexer(
            connection_factory=cf,
            public_ip='192.168.1.1',
            relaying=True
        )
        self.assertIsInstance(cm, protocol.DatagramProtocol)
        self.assertEqual(cm.public_ip, '192.168.1.1')
        self.assertTrue(cm.relaying)
        self.assertIsNone(cm.transport)
        self.assertEqual(len(cm), 0)

    def test_make_connection(self):
        transport = proto_helpers.StringTransportWithDisconnection()
        cm = self._make_cm_with_mocks()
        cm.makeConnection(transport)
        self.assertIs(cm.transport, transport)

    def test_get_nonexistent_connection(self):
        cm = self._make_cm_with_mocks()
        self.assertNotIn(self.addr, cm)
        with self.assertRaises(KeyError):
            con = cm[self.addr]

    def test_set_and_get_new_connection(self):
        cm = self._make_cm_with_mocks()
        mock_connection = mock.Mock(spec_set=connection.RUDPConnection)
        cm[self.addr] = mock_connection
        self.assertIn(self.addr, cm)
        self.assertIs(cm[self.addr], mock_connection)

    def test_set_existent_connection(self):
        pass
<|endoftext|>


<commit_before>import unittest
import time
from tests.common import load_check
from nose.plugins.attrib import attr

# This test is dependent of having a fully open snmpd responding at localhost:161
# with an authentication by the Community String "public"
# This setup should normally be handled by the .travis.yml file, look there if
# you want to see how to run these tests locally

@attr(requires='snmpd')
class TestSNMP(unittest.TestCase):

    def setUp(self):
        self.agentConfig = {
            'version': '0.1',
            'api_key': 'toto'
        }

    def testInit(self):
        self.config = {
                "init_config": {
                    'mibs_folder':'/etc/mibs'
                    }
                }
        # Initialize the check from checks.d
        self.check = load_check('snmp', self.config, self.agentConfig)

        mib_folders = self.check.cmd_generator.snmpEngine.msgAndPduDsp\
                .mibInstrumController.mibBuilder.getMibSources()
        custom_folder_represented = False
        for folder in mib_folders:
            if '/etc/mibs' == folder.fullPath():
                custom_folder_represented = True
                break
        self.assertTrue(custom_folder_represented)

    def test_scalar_SNMPCheck(self):
        self.config = {
                "instances": [{
                    "ip_address": "localhost",
                    "port":161,
                    "community_string": "public",
                    "metrics": [{
                        "OID": "1.3.6.1.2.1.7.1.0", # Counter (needlessly specify the index)
                        "name": "udpDatagrams"
                        },{
                        "OID": "1.3.6.1.2.1.6.10", # Counter
                        "name": "tcpInSegs"
                        },{
                        "MIB": "TCP-MIB",          # Gauge
                        "symbol": "tcpCurrEstab",
                        }]
                    }]
                }

        self.check = load_check('snmp', self.config, self.agentConfig)

        self.check.check(self.config['instances'][0])
        metrics = self.check.get_metrics()

        # Assert that there is only the gauge metric because the counters are
        # used as rate so we don't report them with 1 point
        self.assertEqual(len(metrics), 1)
        self.assertEqual(metrics[0][0], 'snmp.tcpCurrEstab')

        # Sleep for 1 second so the rate interval >=1
        time.sleep(1)
        # Run the check again so we get the rate
        self.check.check(self.config['instances'][0])
        metrics = self.check.get_metrics()

        self.assertEqual(len(metrics) ,3)
        expected_metrics = ['snmp.udpDatagrams', 'snmp.tcpCurrEstab', 'snmp.tcpInSegs']
        for metric in expected_metrics:
            for result in metrics:
                if result[0] == metric:
                    break
            else:
                self.fail("Missing metric: %s" % metric)

        # Service checks
        service_checks = self.check.get_service_checks()
        service_checks = [sc for sc in service_checks if sc['check'].startswith('snmp')]
        service_checks_count = len(service_checks)
        # We run the check twice
        self.assertEquals(service_checks_count, 2, service_checks)
        for sc in service_checks:
            self.assertEquals(sc['status'], self.check.OK, sc)
            self.assertEquals(sc['tags'], ['snmp_device:localhost'], sc)

    def test_table_SNMPCheck(self):
        self.config = {
                "instances": [{
                    "ip_address": "localhost",
                    "port":161,
                    "community_string": "public",
                    "metrics": [{
                        "MIB": "IF-MIB",
                        "table": "ifTable",
                        "symbols": ["ifInOctets", "ifOutOctets"],
                        "metric_tags": [{
                            "tag":"interface",
                            "column":"ifDescr"
                            }, {
                            "tag":"dumbindex",
                            "index":1
                            }]
                        }]
                    }]
                }

        self.check = load_check('snmp', self.config, self.agentConfig)

        self.check.check(self.config['instances'][0])
        # Sleep for 1 second so the rate interval >=1
        time.sleep(1)
        # Run the check again so that we get the rates
        self.check.check(self.config['instances'][0])
        metrics = self.check.get_metrics()
        # nb of metrics depends on the nb of interfaces on the test machine
        # so it's not possible to specify an excat number
        self.assertTrue(len(metrics)>0, "No metrics")
        for metric in metrics:
            self.assertTrue(metric[0] in ['snmp.ifInOctets', 'snmp.ifOutOctets'],
                            metric[0])
            tags = metric[3]['tags']
            # Assert that all the wanted tags are here
            self.assertEquals(len(tags), 3, tags)
            tag_group_expected = ["snmp_device", "dumbindex", "interface"]
            for tag in tags:
                tag_group = tag.split(":")[0]
                self.assertTrue(tag_group in tag_group_expected, tag_group)
                if tag_group == "interface":
                    interface_type = tag.split(":")[1]
                    try:
                        float(interface_type)
                    except ValueError:
                        pass
                    else:
                        self.fail("Tag discovered not pretty printed %s" % interface_type)

        # Service checks
        service_checks = self.check.get_service_checks()
        service_checks = [sc for sc in service_checks if sc['check'].startswith('snmp')]
        service_checks_count = len(service_checks)
        # We run the check twice
        self.assertEquals(service_checks_count, 2, service_checks)
        for sc in service_checks:
            self.assertEquals(sc['status'], self.check.OK, sc)
            self.assertEquals(sc['tags'], ['snmp_device:localhost'], sc)

    def test_network_error(self):
        self.config = {
                "instances": [{
                    "ip_address": "localhost",
                    "port":162,
                    "community_string": "public",
                    "metrics": [{
                        "MIB": "IF-MIB",
                        "table": "ifTable",
                        "symbols": ["ifInOctets", "ifOutOctets"],
                        "metric_tags": [{
                            "tag":"interface",
                            "column":"ifDescr"
                            }, {
                            "tag":"dumbindex",
                            "index":1
                            }]
                        }]
                    }]
                }

        self.check = load_check('snmp', self.config, self.agentConfig)

        # Make it fails faster
        self.check.RETRIES = 0
        self.check.TIMEOUT = 0.5

        # We expect: No SNMP response received before timeout for instance localhost
        with self.assertRaises(Exception):
            self.check.check(self.config['instances'][0])

        # Service checks
        service_checks = self.check.get_service_checks()
        service_checks = [sc for sc in service_checks if sc['check'].startswith('snmp')]
        service_checks_count = len(service_checks)
        self.assertEquals(service_checks_count, 1, service_checks)
        for sc in service_checks:
            self.assertEquals(sc['status'], self.check.CRITICAL, sc)
            self.assertEquals(sc['tags'], ['snmp_device:localhost'], sc)

    def test_invalid_metric(self):
        self.config = {
                "instances": [{
                    "ip_address": "localhost",
                    "port":161,
                    "community_string": "public",
                    "metrics": [{
                        "MIB": "IF-MIB",
                        "table": "ifTable",
                        "symbols": ["ifInOctets", "ifOutOctets"],
                    },{
                        "MIB": "IF-MIB",
                        "table": "noIdeaWhatIAmDoingHere",
                        "symbols": ["ifInOctets", "ifOutOctets"],
                    }]
                }]
            }

        self.check = load_check('snmp', self.config, self.agentConfig)

        # Make it fails faster
        self.check.RETRIES = 0
        self.check.TIMEOUT = 0.5

        # We expect: No symbol IF-MIB::noIdeaWhatIAmDoingHere
        with self.assertRaises(Exception):
            self.check.check(self.config['instances'][0])

        # Service checks
        service_checks = self.check.get_service_checks()
        service_checks = [sc for sc in service_checks if sc['check'].startswith('snmp')]
        service_checks_count = len(service_checks)
        self.assertEquals(service_checks_count, 1, service_checks)
        for sc in service_checks:
            self.assertEquals(sc['status'], self.check.CRITICAL, sc)
            self.assertEquals(sc['tags'], ['snmp_device:localhost'], sc)

if __name__ == "__main__":
    unittest.main()
<commit_msg>Fix SNMP tests for Python 2.6<commit_after>import unittest
import time
from tests.common import load_check
from nose.plugins.attrib import attr

# This test is dependent of having a fully open snmpd responding at localhost:161
# with an authentication by the Community String "public"
# This setup should normally be handled by the .travis.yml file, look there if
# you want to see how to run these tests locally

@attr(requires='snmpd')
class TestSNMP(unittest.TestCase):

    def setUp(self):
        self.agentConfig = {
            'version': '0.1',
            'api_key': 'toto'
        }

    def testInit(self):
        self.config = {
                "init_config": {
                    'mibs_folder':'/etc/mibs'
                    }
                }
        # Initialize the check from checks.d
        self.check = load_check('snmp', self.config, self.agentConfig)

        mib_folders = self.check.cmd_generator.snmpEngine.msgAndPduDsp\
                .mibInstrumController.mibBuilder.getMibSources()
        custom_folder_represented = False
        for folder in mib_folders:
            if '/etc/mibs' == folder.fullPath():
                custom_folder_represented = True
                break
        self.assertTrue(custom_folder_represented)

    def test_scalar_SNMPCheck(self):
        self.config = {
                "instances": [{
                    "ip_address": "localhost",
                    "port":161,
                    "community_string": "public",
                    "metrics": [{
                        "OID": "1.3.6.1.2.1.7.1.0", # Counter (needlessly specify the index)
                        "name": "udpDatagrams"
                        },{
                        "OID": "1.3.6.1.2.1.6.10", # Counter
                        "name": "tcpInSegs"
                        },{
                        "MIB": "TCP-MIB",          # Gauge
                        "symbol": "tcpCurrEstab",
                        }]
                    }]
                }

        self.check = load_check('snmp', self.config, self.agentConfig)

        self.check.check(self.config['instances'][0])
        metrics = self.check.get_metrics()

        # Assert that there is only the gauge metric because the counters are
        # used as rate so we don't report them with 1 point
        self.assertEqual(len(metrics), 1)
        self.assertEqual(metrics[0][0], 'snmp.tcpCurrEstab')

        # Sleep for 1 second so the rate interval >=1
        time.sleep(1)
        # Run the check again so we get the rate
        self.check.check(self.config['instances'][0])
        metrics = self.check.get_metrics()

        self.assertEqual(len(metrics) ,3)
        expected_metrics = ['snmp.udpDatagrams', 'snmp.tcpCurrEstab', 'snmp.tcpInSegs']
        for metric in expected_metrics:
            for result in metrics:
                if result[0] == metric:
                    break
            else:
                self.fail("Missing metric: %s" % metric)

        # Service checks
        service_checks = self.check.get_service_checks()
        service_checks = [sc for sc in service_checks if sc['check'].startswith('snmp')]
        service_checks_count = len(service_checks)
        # We run the check twice
        self.assertEquals(service_checks_count, 2, service_checks)
        for sc in service_checks:
            self.assertEquals(sc['status'], self.check.OK, sc)
            self.assertEquals(sc['tags'], ['snmp_device:localhost'], sc)

    def test_table_SNMPCheck(self):
        self.config = {
                "instances": [{
                    "ip_address": "localhost",
                    "port":161,
                    "community_string": "public",
                    "metrics": [{
                        "MIB": "IF-MIB",
                        "table": "ifTable",
                        "symbols": ["ifInOctets", "ifOutOctets"],
                        "metric_tags": [{
                            "tag":"interface",
                            "column":"ifDescr"
                            }, {
                            "tag":"dumbindex",
                            "index":1
                            }]
                        }]
                    }]
                }

        self.check = load_check('snmp', self.config, self.agentConfig)

        self.check.check(self.config['instances'][0])
        # Sleep for 1 second so the rate interval >=1
        time.sleep(1)
        # Run the check again so that we get the rates
        self.check.check(self.config['instances'][0])
        metrics = self.check.get_metrics()
        # nb of metrics depends on the nb of interfaces on the test machine
        # so it's not possible to specify an excat number
        self.assertTrue(len(metrics)>0, "No metrics")
        for metric in metrics:
            self.assertTrue(metric[0] in ['snmp.ifInOctets', 'snmp.ifOutOctets'],
                            metric[0])
            tags = metric[3]['tags']
            # Assert that all the wanted tags are here
            self.assertEquals(len(tags), 3, tags)
            tag_group_expected = ["snmp_device", "dumbindex", "interface"]
            for tag in tags:
                tag_group = tag.split(":")[0]
                self.assertTrue(tag_group in tag_group_expected, tag_group)
                if tag_group == "interface":
                    interface_type = tag.split(":")[1]
                    try:
                        float(interface_type)
                    except ValueError:
                        pass
                    else:
                        self.fail("Tag discovered not pretty printed %s" % interface_type)

        # Service checks
        service_checks = self.check.get_service_checks()
        service_checks = [sc for sc in service_checks if sc['check'].startswith('snmp')]
        service_checks_count = len(service_checks)
        # We run the check twice
        self.assertEquals(service_checks_count, 2, service_checks)
        for sc in service_checks:
            self.assertEquals(sc['status'], self.check.OK, sc)
            self.assertEquals(sc['tags'], ['snmp_device:localhost'], sc)

    def test_network_error(self):
        self.config = {
                "instances": [{
                    "ip_address": "localhost",
                    "port":162,
                    "community_string": "public",
                    "metrics": [{
                        "MIB": "IF-MIB",
                        "table": "ifTable",
                        "symbols": ["ifInOctets", "ifOutOctets"],
                        "metric_tags": [{
                            "tag":"interface",
                            "column":"ifDescr"
                            }, {
                            "tag":"dumbindex",
                            "index":1
                            }]
                        }]
                    }]
                }

        self.check = load_check('snmp', self.config, self.agentConfig)

        # Make it fails faster
        self.check.RETRIES = 0
        self.check.TIMEOUT = 0.5

        # We expect: No SNMP response received before timeout for instance localhost
        self.assertRaises(Exception, self.check.check, self.config['instances'][0])

        # Service checks
        service_checks = self.check.get_service_checks()
        service_checks = [sc for sc in service_checks if sc['check'].startswith('snmp')]
        service_checks_count = len(service_checks)
        self.assertEquals(service_checks_count, 1, service_checks)
        for sc in service_checks:
            self.assertEquals(sc['status'], self.check.CRITICAL, sc)
            self.assertEquals(sc['tags'], ['snmp_device:localhost'], sc)

    def test_invalid_metric(self):
        self.config = {
                "instances": [{
                    "ip_address": "localhost",
                    "port":161,
                    "community_string": "public",
                    "metrics": [{
                        "MIB": "IF-MIB",
                        "table": "ifTable",
                        "symbols": ["ifInOctets", "ifOutOctets"],
                    },{
                        "MIB": "IF-MIB",
                        "table": "noIdeaWhatIAmDoingHere",
                        "symbols": ["ifInOctets", "ifOutOctets"],
                    }]
                }]
            }

        self.check = load_check('snmp', self.config, self.agentConfig)

        # Make it fails faster
        self.check.RETRIES = 0
        self.check.TIMEOUT = 0.5

        # We expect: No symbol IF-MIB::noIdeaWhatIAmDoingHere
        self.assertRaises(Exception, self.check.check, self.config['instances'][0])

        # Service checks
        service_checks = self.check.get_service_checks()
        service_checks = [sc for sc in service_checks if sc['check'].startswith('snmp')]
        service_checks_count = len(service_checks)
        self.assertEquals(service_checks_count, 1, service_checks)
        for sc in service_checks:
            self.assertEquals(sc['status'], self.check.CRITICAL, sc)
            self.assertEquals(sc['tags'], ['snmp_device:localhost'], sc)

if __name__ == "__main__":
    unittest.main()
<|endoftext|>