<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
# by Al Nikolov <root@toor.fi.eu.org>

import bz2
import ConfigParser
import glob
import httplib
import imp
try:
    import json
except ImportError:
    import simplejson as json
import logging
import os
import pickle
import Queue
import signal
import socket
import StringIO
import subprocess
import sys
import threading
import time
import types
import urllib
import urllib2

__version__ = '1.1.12’  # App version

__FILEABSDIRNAME__ = os.path.dirname(os.path.abspath(__file__))

ini_files = (
    os.path.join('/etc', 'nixstats.ini'),
    os.path.join('/etc', 'nixstats-token.ini'),
    os.path.join(__FILEABSDIRNAME__, 'nixstats.ini'),
    os.path.join(__FILEABSDIRNAME__, 'nixstats-token.ini'),
    os.path.abspath('nixstats.ini'),
    os.path.abspath('nixstats-token.ini'),
)

def info():
    '''
    Return string with info about nixstatsagent:
        - version
        - plugins enabled
        - absolute path to plugin directory
        - server id from configuration file
    '''
    agent = Agent(dry_instance=True)
    plugins_path = agent.config.get('agent', 'plugins')

    plugins_enabled = agent._get_plugins(state='enabled')

    return '\n'.join((
        'Version: %s' % __version__,
        'Plugins enabled: %s' % ', '.join(plugins_enabled),
        'Plugins directory: %s' % plugins_path,
        'Server: %s' % agent.config.get('agent', 'server'),
    ))

def hello(proto='https'):
    user_id = sys.argv[1]
    if len(sys.argv) > 2:
        token_filename = sys.argv[2]
    else:
        token_filename = os.path.join(__FILEABSDIRNAME__, 'nixstats-token.ini')
    if os.path.isfile('/etc/nixstats/token'):
        oldconfigfile = open('/etc/nixstats/token','r')
        server_id = oldconfigfile.readline()
        print "Upgrading from old monitoring agent"
        print "Remove the old agent from the crontab (crontab -e -u nixstats)"
    elif os.path.isfile('/opt/nixstats/nixstats.cfg'):
        oldconfigfile = open('/opt/nixstats/nixstats.cfg')
        lines=oldconfigfile.readlines()
        server_id = lines[1].replace('server=', '').strip()
        print "Upgrading from old python client."
        print "Run :\nchkconfig --del nixstats \nor \nupdate-rc.d -f nixstats remove \nto remove the old service."
    else:
        try:
            hostname = os.uname()[1]
        except AttributeError:
            hostname = socket.getfqdn()
        server_id = urllib2.urlopen(
            proto + '://api.nixstats.com/hello.php',
            data=urllib.urlencode({
                    'user': user_id,
                    'hostname': hostname
            })
        ).read()
    print('Got server_id: %s' % server_id)
    open(token_filename, 'w').\
        write('[DEFAULT]\nuser=%s\nserver=%s\n' % (user_id, server_id))

# def run_agent():
#     Agent().run()

def _plugin_name(plugin):
    if isinstance(plugin, basestring):
        basename = os.path.basename(plugin)
        return os.path.splitext(basename)[0]
    else:
        return plugin.__name__

def test_plugins(plugins=[]):
    '''
    Test specified plugins and print their data output after single check.
    If plugins list is empty test all enabled plugins.
    '''
    agent = Agent(dry_instance=True)
    plugins_path = agent.config.get('agent', 'plugins')
    if plugins_path not in sys.path:
        sys.path.insert(0, plugins_path)

    if not plugins:
        plugins = agent._get_plugins(state='enabled')
        print 'Check all enabled plugins: %s' % ', '.join(plugins)

    for plugin_name in plugins:
        print '%s:' % plugin_name

        try:
            fp, pathname, description = imp.find_module(plugin_name)
        except Exception as e:
            print 'Find error:', e
            continue

        try:
            module = imp.load_module(plugin_name, fp, pathname, description)
        except Exception as e:
            print 'Load error:', e
            continue
        finally:
            # Since we may exit via an exception, close fp explicitly.
            if fp:
                fp.close()

        try:
            payload = module.Plugin().run(agent.config)
            print json.dumps(payload, indent=4, sort_keys=True)
        except Exception as e:
            print 'Execution error:', e

class Agent:
    execute = Queue.Queue()
    metrics = Queue.Queue()
    data = Queue.Queue()
    shutdown = False

    def __init__(self, dry_instance=False):
        '''
        Initialize internal strictures
        '''
        self._config_init()

        # Cache for plugins so they can store values related to previous checks
        self.plugins_cache = {}

        if dry_instance:
            return

        self._logging_init()
        self._plugins_init()
        self._data_worker_init()
        self._dump_config()

    def _config_init(self):
        '''
        Initialize configuration object
        '''
        defaults = {
            'max_data_span': 60,
            'max_data_age': 60 * 10,
            'logging_level': logging.WARNING,
            'threads': 100,
            'ttl': 60,
            'interval': 60,
            'plugins': os.path.join(__FILEABSDIRNAME__, 'plugins'),
            'enabled': 'no',
            'subprocess': 'no',
            'user': '',
            'server': '',
            'api_host': 'api.nixstats.com',
            'api_path': '/v2/server/poll',
        }
        sections = [
            'agent',
            'execution',
            'data',
        ]
        config = ConfigParser.RawConfigParser(defaults)
        config.read(ini_files)
        self.config = config
        for section in sections:
            self._config_section_create(section)
            if section is 'data' or section is 'agent':
                self.config.set(section, 'interval', 1)

    def _config_section_create(self, section):
        '''
        Create an addition section in the configuration object
        if it's not exists
        '''
        if not self.config.has_section(section):
            self.config.add_section(section)

    def _logging_init(self):
        '''
        Initialize logging faculty
        '''
        level = self.config.getint('agent', 'logging_level')
        logging.basicConfig(level=level)
        logging.info('Agent logging_level %i', level)

    def _plugins_init(self):
        '''
        Discover the plugins
        '''
        logging.info('_plugins_init')
        plugins_path = self.config.get('agent', 'plugins')
        filenames = glob.glob(os.path.join(plugins_path, '*.py'))
        if plugins_path not in sys.path:
            sys.path.insert(0, plugins_path)
        self.schedule = {}
        for filename in filenames:
            name = _plugin_name(filename)
            if name == 'plugins':
                continue
            self._config_section_create(name)
            if self.config.getboolean(name, 'enabled'):
                if self.config.getboolean(name, 'subprocess'):
                    self.schedule[filename] = 0
                else:
                    fp, pathname, description = imp.find_module(name)
                    try:
                        module = imp.load_module(name, fp, pathname, description)
                    except Exception:
                        module = None
                    finally:
                        # Since we may exit via an exception, close fp explicitly.
                        if fp:
                            fp.close()
                    if module:
                        self.schedule[module] = 0
                    else:
                        logging.error('import_plugin:%s:%s', name, sys.exc_type)

    def _subprocess_execution(self, task):
        '''
        Execute /task/ in a subprocess
        '''
        process = subprocess.Popen((sys.executable, task),
            stdout=subprocess.PIPE, stderr=subprocess.PIPE,
            universal_newlines=True)
        logging.info('%s:process:%i', threading.currentThread(), process.pid)
        interval = self.config.getint('execution', 'interval')
        name = _plugin_name(task)
        ttl = self.config.getint(name, 'ttl')
        ticks = ttl / interval or 1
        process.poll()
        while process.returncode is None and ticks > 0:
            logging.info('%s:tick:%i', threading.currentThread(), ticks)
            time.sleep(interval)
            ticks -= 1
            process.poll()
        if process.returncode is None:
            logging.error('%s:kill:%i', threading.currentThread(), process.pid)
            os.kill(process.pid, signal.SIGTERM)
        stdout, stderr = process.communicate()
        if process.returncode != 0 or stderr:
            logging.error('%s:%s:%s:%s', threading.currentThread(),
                task, process.returncode, stderr)
        if stdout:
            ret = pickle.loads(stdout)
        else:
            ret = None
        return ret

    def _execution(self):
        '''
        Take queued execution requests, execute plugins and queue the results
        '''
        while True:
            if self.shutdown:
                logging.info('%s:shutdown', threading.currentThread())
                break
            logging.info('%s:exec_queue:%i', threading.currentThread(), self.execute.qsize())
            try:
                task = self.execute.get_nowait()
            except Queue.Empty:
                break
            logging.info('%s:task:%s', threading.currentThread(), task)
            name = _plugin_name(task)
            ts = time.time()
            if isinstance(task, basestring):
                payload = self._subprocess_execution(task)
            else:
                try:
                    # Setup cache for plugin instance
                    # if name not in self.plugins_cache.iterkeys():
                    #     self.plugins_cache[name] = []
                    self.plugins_cache.update({
                        name: self.plugins_cache.get(name, [])
                    })

                    plugin = task.Plugin(agent_cache=self.plugins_cache[name])
                    payload = plugin.run(self.config)
                except Exception:
                    logging.exception('plugin_exception')
                    payload = {'exception': str(sys.exc_info()[0])}
            self.metrics.put({
                'ts': ts,
                'task': task,
                'name': name,
                'payload': payload,
            })
            self.hire.release()

    def _data(self):
        '''
        Take and collect data, send and clean if needed
        '''
        logging.info('%s', threading.currentThread())
        collection = []
        while True:
            if self.shutdown:
                logging.info('%s:shutdown', threading.currentThread())
                break
            logging.info('%s:data_queue:%i:collection:%i',
                threading.currentThread(), self.data.qsize(), len(collection))
            while self.data.qsize():
                collection.append(self.data.get_nowait())
            if collection:
                first_ts = min((e['ts'] for e in collection))
                last_ts = max((e['ts'] for e in collection))
                now = time.time()
                max_age = self.config.getint('agent', 'max_data_age')
                max_span = self.config.getint('agent', 'max_data_span')
                send = False
                server = self.config.get('agent', 'server')
                user = self.config.get('agent', 'user')
                if last_ts - first_ts > max_span:
                    logging.info('Max data span')
                    send = True
                elif now - first_ts > max_age:
                    logging.info('Max data age')
                    send = True
                if send:
                    headers = {
                        "Content-type": "application/json",
                        "Authorization":
                            'ApiKey %s:%s' % (
                                self.config.get('agent', 'user'),
                                self.config.get('agent', 'server'),
                            ),
                    }
                    clean = False
                    logging.info('collection:%s', collection)
                    if not server and not user:
                        logging.info('Empty server/user, but need to send: %s',
                            json.dumps(collection))
                        clean = True
                    else:
                        connection = httplib.HTTPSConnection(self.config.get('data', 'api_host'))
                        connection.request('PUT', '%s?version=%s' % (self.config.get('data', 'api_path'), __version__),
                                bz2.compress(str(json.dumps(collection)) + "\n"),
                                headers=headers)
                        response = connection.getresponse()
                        logging.info('%s', response)
                        connection.close()
                        if response.status == 200:
                            clean = True
                    if clean:
                        collection = []
            time.sleep(self.config.getint('data', 'interval'))

    def _data_worker_init(self):
        '''
        Initialize data worker thread
        '''
        logging.info('_data_worker_init')
        threading.Thread(target=self._data).start()

    def _dump_config(self):
        '''
        Dumps configuration object
        '''
        buf = StringIO.StringIO()
        self.config.write(buf)
        logging.info('Config: %s', buf.getvalue())

    def _get_plugins(self, state='enabled'):
        '''
        Return list with plugins names
        '''
        plugins_path = self.config.get('agent', 'plugins')
        plugins = []
        for filename in glob.glob(os.path.join(plugins_path, '*.py')):
            plugin_name = _plugin_name(filename)
            if plugin_name == 'plugins':
                continue
            self._config_section_create(plugin_name)

            if state == 'enabled':
                if self.config.getboolean(plugin_name, 'enabled'):
                    plugins.append(plugin_name)
            elif state == 'disabled':
                if not self.config.getboolean(plugin_name, 'enabled'):
                    plugins.append(plugin_name)

        return plugins

    def run(self):
        '''
        Start all the worker threads
        '''
        logging.info('Agent main loop')
        interval = self.config.getint('agent', 'interval')
        self.hire = threading.Semaphore(
            self.config.getint('execution', 'threads'))
        try:
            while True:
                now = time.time()
                logging.info('%i threads', threading.activeCount())
                while self.metrics.qsize():
                    metrics = self.metrics.get_nowait()
                    name = metrics['name']
                    logging.info('metrics:%s', name)
                    plugin = metrics.get('task')
                    if plugin:
                        self.schedule[plugin] = \
                            now + self.config.getint(name, 'interval')
                        if isinstance(plugin, types.ModuleType):
                            metrics['task'] = plugin.__file__
                    self.data.put(metrics)
                execute = [
                    what
                    for what, when in self.schedule.items()
                    if when <= now
                ]
                for name in execute:
                    logging.info('scheduling:%s', name)
                    del self.schedule[name]
                    self.execute.put(name)
                    if self.hire.acquire(False):
                        thread = threading.Thread(target=self._execution)
                        thread.start()
                        logging.info('new_execution_worker_thread:%s', thread)
                    else:
                        logging.warning('threads_capped')
                        self.metrics.put({
                            'ts': now,
                            'name': 'agent_internal',
                            'payload': {
                                'threads_capping':
                                    self.config.getint('execution', 'threads')}
                        })

                time.sleep(interval)
        except KeyboardInterrupt:
            logging.warning(sys.exc_info()[0])
            while True:
                wait_for = [thread
                    for thread in threading.enumerate()
                        if not thread.isDaemon() and
                           not isinstance(thread, threading._MainThread)]
                logging.info('Shutdown, waiting for %i threads to exit', len(wait_for))
                logging.info('Remaining threads: %s', threading.enumerate())
                if len(wait_for) == 0:
                    sys.exit(0)
                self.shutdown = True
                time.sleep(interval)

def main():
    if len(sys.argv) > 1:
        if sys.argv[1].startswith('--'):
            sys.argv[1] = sys.argv[1][2:]

        if sys.argv[1] == 'info':
            print info()
            sys.exit()
        elif sys.argv[1] == 'version':
            print __version__
            sys.exit()
        elif sys.argv[1] == 'hello':
            del sys.argv[1]
            hello()
        elif sys.argv[1] == 'insecure-hello':
            del sys.argv[1]
            hello(proto='http')
        elif sys.argv[1] == 'test':
            test_plugins(sys.argv[2:])
            sys.exit()
        else:
            print >>sys.stderr, 'Invalid option:', sys.argv[1]
            sys.exit(1)
    else:
        Agent().run()

if __name__ == '__main__':
    main()
<commit_msg>Update nixstatsagent.py<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
# by Al Nikolov <root@toor.fi.eu.org>

import bz2
import ConfigParser
import glob
import httplib
import imp
try:
    import json
except ImportError:
    import simplejson as json
import logging
import os
import pickle
import Queue
import signal
import socket
import StringIO
import subprocess
import sys
import threading
import time
import types
import urllib
import urllib2

__version__ = '1.1.12'  # App version

__FILEABSDIRNAME__ = os.path.dirname(os.path.abspath(__file__))

ini_files = (
    os.path.join('/etc', 'nixstats.ini'),
    os.path.join('/etc', 'nixstats-token.ini'),
    os.path.join(__FILEABSDIRNAME__, 'nixstats.ini'),
    os.path.join(__FILEABSDIRNAME__, 'nixstats-token.ini'),
    os.path.abspath('nixstats.ini'),
    os.path.abspath('nixstats-token.ini'),
)

def info():
    '''
    Return string with info about nixstatsagent:
        - version
        - plugins enabled
        - absolute path to plugin directory
        - server id from configuration file
    '''
    agent = Agent(dry_instance=True)
    plugins_path = agent.config.get('agent', 'plugins')

    plugins_enabled = agent._get_plugins(state='enabled')

    return '\n'.join((
        'Version: %s' % __version__,
        'Plugins enabled: %s' % ', '.join(plugins_enabled),
        'Plugins directory: %s' % plugins_path,
        'Server: %s' % agent.config.get('agent', 'server'),
    ))

def hello(proto='https'):
    user_id = sys.argv[1]
    if len(sys.argv) > 2:
        token_filename = sys.argv[2]
    else:
        token_filename = os.path.join(__FILEABSDIRNAME__, 'nixstats-token.ini')
    if os.path.isfile('/etc/nixstats/token'):
        oldconfigfile = open('/etc/nixstats/token','r')
        server_id = oldconfigfile.readline()
        print "Upgrading from old monitoring agent"
        print "Remove the old agent from the crontab (crontab -e -u nixstats)"
    elif os.path.isfile('/opt/nixstats/nixstats.cfg'):
        oldconfigfile = open('/opt/nixstats/nixstats.cfg')
        lines=oldconfigfile.readlines()
        server_id = lines[1].replace('server=', '').strip()
        print "Upgrading from old python client."
        print "Run :\nchkconfig --del nixstats \nor \nupdate-rc.d -f nixstats remove \nto remove the old service."
    else:
        try:
            hostname = os.uname()[1]
        except AttributeError:
            hostname = socket.getfqdn()
        server_id = urllib2.urlopen(
            proto + '://api.nixstats.com/hello.php',
            data=urllib.urlencode({
                    'user': user_id,
                    'hostname': hostname
            })
        ).read()
    print('Got server_id: %s' % server_id)
    open(token_filename, 'w').\
        write('[DEFAULT]\nuser=%s\nserver=%s\n' % (user_id, server_id))

# def run_agent():
#     Agent().run()

def _plugin_name(plugin):
    if isinstance(plugin, basestring):
        basename = os.path.basename(plugin)
        return os.path.splitext(basename)[0]
    else:
        return plugin.__name__

def test_plugins(plugins=[]):
    '''
    Test specified plugins and print their data output after single check.
    If plugins list is empty test all enabled plugins.
    '''
    agent = Agent(dry_instance=True)
    plugins_path = agent.config.get('agent', 'plugins')
    if plugins_path not in sys.path:
        sys.path.insert(0, plugins_path)

    if not plugins:
        plugins = agent._get_plugins(state='enabled')
        print 'Check all enabled plugins: %s' % ', '.join(plugins)

    for plugin_name in plugins:
        print '%s:' % plugin_name

        try:
            fp, pathname, description = imp.find_module(plugin_name)
        except Exception as e:
            print 'Find error:', e
            continue

        try:
            module = imp.load_module(plugin_name, fp, pathname, description)
        except Exception as e:
            print 'Load error:', e
            continue
        finally:
            # Since we may exit via an exception, close fp explicitly.
            if fp:
                fp.close()

        try:
            payload = module.Plugin().run(agent.config)
            print json.dumps(payload, indent=4, sort_keys=True)
        except Exception as e:
            print 'Execution error:', e

class Agent:
    execute = Queue.Queue()
    metrics = Queue.Queue()
    data = Queue.Queue()
    shutdown = False

    def __init__(self, dry_instance=False):
        '''
        Initialize internal strictures
        '''
        self._config_init()

        # Cache for plugins so they can store values related to previous checks
        self.plugins_cache = {}

        if dry_instance:
            return

        self._logging_init()
        self._plugins_init()
        self._data_worker_init()
        self._dump_config()

    def _config_init(self):
        '''
        Initialize configuration object
        '''
        defaults = {
            'max_data_span': 60,
            'max_data_age': 60 * 10,
            'logging_level': logging.WARNING,
            'threads': 100,
            'ttl': 60,
            'interval': 60,
            'plugins': os.path.join(__FILEABSDIRNAME__, 'plugins'),
            'enabled': 'no',
            'subprocess': 'no',
            'user': '',
            'server': '',
            'api_host': 'api.nixstats.com',
            'api_path': '/v2/server/poll',
        }
        sections = [
            'agent',
            'execution',
            'data',
        ]
        config = ConfigParser.RawConfigParser(defaults)
        config.read(ini_files)
        self.config = config
        for section in sections:
            self._config_section_create(section)
            if section is 'data' or section is 'agent':
                self.config.set(section, 'interval', 1)

    def _config_section_create(self, section):
        '''
        Create an addition section in the configuration object
        if it's not exists
        '''
        if not self.config.has_section(section):
            self.config.add_section(section)

    def _logging_init(self):
        '''
        Initialize logging faculty
        '''
        level = self.config.getint('agent', 'logging_level')
        logging.basicConfig(level=level)
        logging.info('Agent logging_level %i', level)

    def _plugins_init(self):
        '''
        Discover the plugins
        '''
        logging.info('_plugins_init')
        plugins_path = self.config.get('agent', 'plugins')
        filenames = glob.glob(os.path.join(plugins_path, '*.py'))
        if plugins_path not in sys.path:
            sys.path.insert(0, plugins_path)
        self.schedule = {}
        for filename in filenames:
            name = _plugin_name(filename)
            if name == 'plugins':
                continue
            self._config_section_create(name)
            if self.config.getboolean(name, 'enabled'):
                if self.config.getboolean(name, 'subprocess'):
                    self.schedule[filename] = 0
                else:
                    fp, pathname, description = imp.find_module(name)
                    try:
                        module = imp.load_module(name, fp, pathname, description)
                    except Exception:
                        module = None
                    finally:
                        # Since we may exit via an exception, close fp explicitly.
                        if fp:
                            fp.close()
                    if module:
                        self.schedule[module] = 0
                    else:
                        logging.error('import_plugin:%s:%s', name, sys.exc_type)

    def _subprocess_execution(self, task):
        '''
        Execute /task/ in a subprocess
        '''
        process = subprocess.Popen((sys.executable, task),
            stdout=subprocess.PIPE, stderr=subprocess.PIPE,
            universal_newlines=True)
        logging.info('%s:process:%i', threading.currentThread(), process.pid)
        interval = self.config.getint('execution', 'interval')
        name = _plugin_name(task)
        ttl = self.config.getint(name, 'ttl')
        ticks = ttl / interval or 1
        process.poll()
        while process.returncode is None and ticks > 0:
            logging.info('%s:tick:%i', threading.currentThread(), ticks)
            time.sleep(interval)
            ticks -= 1
            process.poll()
        if process.returncode is None:
            logging.error('%s:kill:%i', threading.currentThread(), process.pid)
            os.kill(process.pid, signal.SIGTERM)
        stdout, stderr = process.communicate()
        if process.returncode != 0 or stderr:
            logging.error('%s:%s:%s:%s', threading.currentThread(),
                task, process.returncode, stderr)
        if stdout:
            ret = pickle.loads(stdout)
        else:
            ret = None
        return ret

    def _execution(self):
        '''
        Take queued execution requests, execute plugins and queue the results
        '''
        while True:
            if self.shutdown:
                logging.info('%s:shutdown', threading.currentThread())
                break
            logging.info('%s:exec_queue:%i', threading.currentThread(), self.execute.qsize())
            try:
                task = self.execute.get_nowait()
            except Queue.Empty:
                break
            logging.info('%s:task:%s', threading.currentThread(), task)
            name = _plugin_name(task)
            ts = time.time()
            if isinstance(task, basestring):
                payload = self._subprocess_execution(task)
            else:
                try:
                    # Setup cache for plugin instance
                    # if name not in self.plugins_cache.iterkeys():
                    #     self.plugins_cache[name] = []
                    self.plugins_cache.update({
                        name: self.plugins_cache.get(name, [])
                    })

                    plugin = task.Plugin(agent_cache=self.plugins_cache[name])
                    payload = plugin.run(self.config)
                except Exception:
                    logging.exception('plugin_exception')
                    payload = {'exception': str(sys.exc_info()[0])}
            self.metrics.put({
                'ts': ts,
                'task': task,
                'name': name,
                'payload': payload,
            })
            self.hire.release()

    def _data(self):
        '''
        Take and collect data, send and clean if needed
        '''
        logging.info('%s', threading.currentThread())
        collection = []
        while True:
            if self.shutdown:
                logging.info('%s:shutdown', threading.currentThread())
                break
            logging.info('%s:data_queue:%i:collection:%i',
                threading.currentThread(), self.data.qsize(), len(collection))
            while self.data.qsize():
                collection.append(self.data.get_nowait())
            if collection:
                first_ts = min((e['ts'] for e in collection))
                last_ts = max((e['ts'] for e in collection))
                now = time.time()
                max_age = self.config.getint('agent', 'max_data_age')
                max_span = self.config.getint('agent', 'max_data_span')
                send = False
                server = self.config.get('agent', 'server')
                user = self.config.get('agent', 'user')
                if last_ts - first_ts > max_span:
                    logging.info('Max data span')
                    send = True
                elif now - first_ts > max_age:
                    logging.info('Max data age')
                    send = True
                if send:
                    headers = {
                        "Content-type": "application/json",
                        "Authorization":
                            'ApiKey %s:%s' % (
                                self.config.get('agent', 'user'),
                                self.config.get('agent', 'server'),
                            ),
                    }
                    clean = False
                    logging.info('collection:%s', collection)
                    if not server and not user:
                        logging.info('Empty server/user, but need to send: %s',
                            json.dumps(collection))
                        clean = True
                    else:
                        connection = httplib.HTTPSConnection(self.config.get('data', 'api_host'))
                        connection.request('PUT', '%s?version=%s' % (self.config.get('data', 'api_path'), __version__),
                                bz2.compress(str(json.dumps(collection)) + "\n"),
                                headers=headers)
                        response = connection.getresponse()
                        logging.info('%s', response)
                        connection.close()
                        if response.status == 200:
                            clean = True
                    if clean:
                        collection = []
            time.sleep(self.config.getint('data', 'interval'))

    def _data_worker_init(self):
        '''
        Initialize data worker thread
        '''
        logging.info('_data_worker_init')
        threading.Thread(target=self._data).start()

    def _dump_config(self):
        '''
        Dumps configuration object
        '''
        buf = StringIO.StringIO()
        self.config.write(buf)
        logging.info('Config: %s', buf.getvalue())

    def _get_plugins(self, state='enabled'):
        '''
        Return list with plugins names
        '''
        plugins_path = self.config.get('agent', 'plugins')
        plugins = []
        for filename in glob.glob(os.path.join(plugins_path, '*.py')):
            plugin_name = _plugin_name(filename)
            if plugin_name == 'plugins':
                continue
            self._config_section_create(plugin_name)

            if state == 'enabled':
                if self.config.getboolean(plugin_name, 'enabled'):
                    plugins.append(plugin_name)
            elif state == 'disabled':
                if not self.config.getboolean(plugin_name, 'enabled'):
                    plugins.append(plugin_name)

        return plugins

    def run(self):
        '''
        Start all the worker threads
        '''
        logging.info('Agent main loop')
        interval = self.config.getint('agent', 'interval')
        self.hire = threading.Semaphore(
            self.config.getint('execution', 'threads'))
        try:
            while True:
                now = time.time()
                logging.info('%i threads', threading.activeCount())
                while self.metrics.qsize():
                    metrics = self.metrics.get_nowait()
                    name = metrics['name']
                    logging.info('metrics:%s', name)
                    plugin = metrics.get('task')
                    if plugin:
                        self.schedule[plugin] = \
                            now + self.config.getint(name, 'interval')
                        if isinstance(plugin, types.ModuleType):
                            metrics['task'] = plugin.__file__
                    self.data.put(metrics)
                execute = [
                    what
                    for what, when in self.schedule.items()
                    if when <= now
                ]
                for name in execute:
                    logging.info('scheduling:%s', name)
                    del self.schedule[name]
                    self.execute.put(name)
                    if self.hire.acquire(False):
                        thread = threading.Thread(target=self._execution)
                        thread.start()
                        logging.info('new_execution_worker_thread:%s', thread)
                    else:
                        logging.warning('threads_capped')
                        self.metrics.put({
                            'ts': now,
                            'name': 'agent_internal',
                            'payload': {
                                'threads_capping':
                                    self.config.getint('execution', 'threads')}
                        })

                time.sleep(interval)
        except KeyboardInterrupt:
            logging.warning(sys.exc_info()[0])
            while True:
                wait_for = [thread
                    for thread in threading.enumerate()
                        if not thread.isDaemon() and
                           not isinstance(thread, threading._MainThread)]
                logging.info('Shutdown, waiting for %i threads to exit', len(wait_for))
                logging.info('Remaining threads: %s', threading.enumerate())
                if len(wait_for) == 0:
                    sys.exit(0)
                self.shutdown = True
                time.sleep(interval)

def main():
    if len(sys.argv) > 1:
        if sys.argv[1].startswith('--'):
            sys.argv[1] = sys.argv[1][2:]

        if sys.argv[1] == 'info':
            print info()
            sys.exit()
        elif sys.argv[1] == 'version':
            print __version__
            sys.exit()
        elif sys.argv[1] == 'hello':
            del sys.argv[1]
            hello()
        elif sys.argv[1] == 'insecure-hello':
            del sys.argv[1]
            hello(proto='http')
        elif sys.argv[1] == 'test':
            test_plugins(sys.argv[2:])
            sys.exit()
        else:
            print >>sys.stderr, 'Invalid option:', sys.argv[1]
            sys.exit(1)
    else:
        Agent().run()

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before>#!/usr/bin/env python
# encoding: utf-8

"""

This file provides classes to use with pygame.

Font class requires font files, located under ./fonts
Mouseresponse class requires cursor files, located under ./cursors

Copyright 2015 Takao Noguchi (tkngch@runbox.com)

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

"""

import os
import pygame

class Button:

    def __init__(self, font, value='Click Me!',
                 boxsize=(400, 150), screen=None):
        self.font = font
        self.value = value
        self.boxsize = boxsize
        self.colour = (0, 0, 0)
        self.bgColour = (200, 200, 200)
        if screen is None:
            self.screen = pygame.display.get_surface()
        else:
            self.screen = screen
        self.center = (300, 200)

    def render(self):
        self.surface = pygame.Surface(self.boxsize)
        self.rect = pygame.Rect(self.surface.get_rect())
        self.rect.center = self.center

        rendered = self.font.render(self.value, 1, self.colour)
        renderedRect = rendered.get_rect()
        renderedRect.center = [s / 2 for s in self.rect.size]

        self.surface.fill(self.bgColour)
        self.surface.blit(rendered, renderedRect)

    def show(self):
        self.render()
        self.screen.blit(self.surface, self.rect)

    def clear(self):
        self.surface.fill((255, 255, 255))
        self.screen.blit(self.surface, self.rect)

class Font:

    def __init__(self, fontname='DejaVuSans', fontsize=25):
        self.fontsize = fontsize
        if not pygame.font.get_init():
            pygame.font.init()

        fontfile = 'fonts/%s.ttf' % fontname
        if os.path.exists(fontfile):
            self.font = pygame.font.Font(fontfile, fontsize)
        elif os.path.exists(fontfile):
            self.font = pygame.font.Font(fontfile, fontsize)
        else:
            raise Exception('Font file does not exist')

    def get(self):
        return self.font

class Image:

    def __init__(self, surface=None, imagefile=None,
                 scale=1, edge=True, edgewidth=1):
        if surface is not None:
            rawImage = surface
        elif imagefile is not None:
            rawImage = pygame.image.load(imagefile)

        if scale is 1:
            image = rawImage
        else:
            imageRect = rawImage.get_rect()
            size = (int(imageRect.size[0] * scale),
                    int(imageRect.size[1] * scale))
            image = pygame.transform.scale(rawImage, size)

        if edge:
            self.imageSurface = self.blackEdge(image, edgewidth)
        else:
            self.imageSurface = image

    def get(self):
        return self.imageSurface

    def blackEdge(self, surface, edgewidth):
        rect = surface.get_rect()

        background = pygame.Surface((rect.width + (edgewidth * 2),
                                     rect.height + (edgewidth * 2)))
        background.fill((0, 0, 0))
        background.get_rect().center = rect.center
        background.blit(surface, pygame.Rect((edgewidth, edgewidth,
                                              rect.width, rect.height)))

        return background

class Keyresponse:

    def __init__(self, acceptedKeys):
        self.acceptedKeys = acceptedKeys

    def wait(self):
        starttime = pygame.time.get_ticks()
        waiting = True
        while waiting:
            key = self.accept_key()
            if key in self.acceptedKeys:
                waiting = False

        keyname = pygame.key.name(key)
        responseTime = pygame.time.get_ticks() - starttime
        return keyname, responseTime

    def accept_key(self):
        """
        accept_key(self)

        returns a pressed key
        """
        pygame.event.clear()
        running = True

        while running:
            for event in pygame.event.get():
                if event.type == pygame.KEYDOWN:
                    targetKey = event.key
                    running = False

        return targetKey

class Mouseresponse:

    def __init__(self, rectangles):
        # rectangles are tuple of pygame rect
        self.rectangles = rectangles

        self.handcursor = pygame.cursors.load_xbm(
                            "cursors/handcursor.xbm",
                            "cursors/handcursor_mask.xbm")

    def wait(self):
        pygame.mouse.set_visible(True)
        self.showSystemcursor()

        pygame.event.clear()
        self.waitRelease()

        starttime = pygame.time.get_ticks()

        handcursor = 0
        running = True
        while running:

            collision = self.checkCollision()
            if ((collision in range(len(self.rectangles))) and
               (handcursor is 0)):

                handcursor = 1
                self.showHandcursor()

            elif ((collision not in range(len(self.rectangles))) and
                  (handcursor is 1)):

                handcursor = 0
                self.showSystemcursor()

            for event in pygame.event.get():
                if event.type == pygame.MOUSEBUTTONDOWN:
                    buttons = pygame.mouse.get_pressed()
                    if buttons[0]:
                        selected = self.checkCollision()
                        if selected is not None:
                            running = False

        responseTime = pygame.time.get_ticks() - starttime

        self.waitRelease()
        self.showSystemcursor()

        return selected, responseTime

    def showHandcursor(self):
        pygame.mouse.set_cursor(*self.handcursor)

    def showSystemcursor(self):
        pygame.mouse.set_cursor(*pygame.cursors.arrow)

    def waitRelease(self):
        while any(pygame.mouse.get_pressed()):
            pygame.event.pump()

    def checkCollision(self):
        mousePosition = pygame.mouse.get_pos()

        selected = None
        for i in range(len(self.rectangles)):
            if self.rectangles[i].collidepoint(mousePosition):
                selected = i

        return selected

class Text:

    def __init__(self, text, screen, font=Font().get(), colour=(0, 0, 0)):
        self.text = text
        self.screen = screen
        self.font = font
        self.colour = colour

    def write(self, position=(0, 0)):
        target = self.font.render(self.text, 1, self.colour)
        targetPosition = target.get_rect()
        targetPosition.centerx = position[0]
        targetPosition.centery = position[1]
        self.screen.blit(target, targetPosition)

        return target, targetPosition

    def wrapWrite(self, rect, lineSpace=1):
        """
        """
        space = self.font.render(" ", 1, self.colour)
        spaceSize = space.get_size()

        leftTop = [rect.left, rect.top]

        for sentence in self.text.split("\n"):
            for word in sentence.split(" "):
                word += " "

                target = self.font.render(word, 1, self.colour)
                targetSize = target.get_size()

                if (leftTop[0] + targetSize[0] > rect.right):
                    # new line
                    leftTop[0] = rect.left
                    leftTop[1] += spaceSize[1] * lineSpace

                self.screen.blit(target, leftTop)

                leftTop[0] += targetSize[0]

            leftTop[0] = rect.left
            leftTop[1] += spaceSize[1] * lineSpace
<commit_msg>hashbang fixed<commit_after>#!/usr/bin/env python2
# encoding: utf-8

"""

This file provides classes to use with pygame.

Font class requires font files, located under ./fonts
Mouseresponse class requires cursor files, located under ./cursors

Copyright 2015 Takao Noguchi (tkngch@runbox.com)

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

"""

import os
import pygame

class Button:

    def __init__(self, font, value='Click Me!',
                 boxsize=(400, 150), screen=None):
        self.font = font
        self.value = value
        self.boxsize = boxsize
        self.colour = (0, 0, 0)
        self.bgColour = (200, 200, 200)
        if screen is None:
            self.screen = pygame.display.get_surface()
        else:
            self.screen = screen
        self.center = (300, 200)

    def render(self):
        self.surface = pygame.Surface(self.boxsize)
        self.rect = pygame.Rect(self.surface.get_rect())
        self.rect.center = self.center

        rendered = self.font.render(self.value, 1, self.colour)
        renderedRect = rendered.get_rect()
        renderedRect.center = [s / 2 for s in self.rect.size]

        self.surface.fill(self.bgColour)
        self.surface.blit(rendered, renderedRect)

    def show(self):
        self.render()
        self.screen.blit(self.surface, self.rect)

    def clear(self):
        self.surface.fill((255, 255, 255))
        self.screen.blit(self.surface, self.rect)

class Font:

    def __init__(self, fontname='DejaVuSans', fontsize=25):
        self.fontsize = fontsize
        if not pygame.font.get_init():
            pygame.font.init()

        fontfile = 'fonts/%s.ttf' % fontname
        if os.path.exists(fontfile):
            self.font = pygame.font.Font(fontfile, fontsize)
        elif os.path.exists(fontfile):
            self.font = pygame.font.Font(fontfile, fontsize)
        else:
            raise Exception('Font file does not exist')

    def get(self):
        return self.font

class Image:

    def __init__(self, surface=None, imagefile=None,
                 scale=1, edge=True, edgewidth=1):
        if surface is not None:
            rawImage = surface
        elif imagefile is not None:
            rawImage = pygame.image.load(imagefile)

        if scale is 1:
            image = rawImage
        else:
            imageRect = rawImage.get_rect()
            size = (int(imageRect.size[0] * scale),
                    int(imageRect.size[1] * scale))
            image = pygame.transform.scale(rawImage, size)

        if edge:
            self.imageSurface = self.blackEdge(image, edgewidth)
        else:
            self.imageSurface = image

    def get(self):
        return self.imageSurface

    def blackEdge(self, surface, edgewidth):
        rect = surface.get_rect()

        background = pygame.Surface((rect.width + (edgewidth * 2),
                                     rect.height + (edgewidth * 2)))
        background.fill((0, 0, 0))
        background.get_rect().center = rect.center
        background.blit(surface, pygame.Rect((edgewidth, edgewidth,
                                              rect.width, rect.height)))

        return background

class Keyresponse:

    def __init__(self, acceptedKeys):
        self.acceptedKeys = acceptedKeys

    def wait(self):
        starttime = pygame.time.get_ticks()
        waiting = True
        while waiting:
            key = self.accept_key()
            if key in self.acceptedKeys:
                waiting = False

        keyname = pygame.key.name(key)
        responseTime = pygame.time.get_ticks() - starttime
        return keyname, responseTime

    def accept_key(self):
        """
        accept_key(self)

        returns a pressed key
        """
        pygame.event.clear()
        running = True

        while running:
            for event in pygame.event.get():
                if event.type == pygame.KEYDOWN:
                    targetKey = event.key
                    running = False

        return targetKey

class Mouseresponse:

    def __init__(self, rectangles):
        # rectangles are tuple of pygame rect
        self.rectangles = rectangles

        self.handcursor = pygame.cursors.load_xbm(
                            "cursors/handcursor.xbm",
                            "cursors/handcursor_mask.xbm")

    def wait(self):
        pygame.mouse.set_visible(True)
        self.showSystemcursor()

        pygame.event.clear()
        self.waitRelease()

        starttime = pygame.time.get_ticks()

        handcursor = 0
        running = True
        while running:

            collision = self.checkCollision()
            if ((collision in range(len(self.rectangles))) and
               (handcursor is 0)):

                handcursor = 1
                self.showHandcursor()

            elif ((collision not in range(len(self.rectangles))) and
                  (handcursor is 1)):

                handcursor = 0
                self.showSystemcursor()

            for event in pygame.event.get():
                if event.type == pygame.MOUSEBUTTONDOWN:
                    buttons = pygame.mouse.get_pressed()
                    if buttons[0]:
                        selected = self.checkCollision()
                        if selected is not None:
                            running = False

        responseTime = pygame.time.get_ticks() - starttime

        self.waitRelease()
        self.showSystemcursor()

        return selected, responseTime

    def showHandcursor(self):
        pygame.mouse.set_cursor(*self.handcursor)

    def showSystemcursor(self):
        pygame.mouse.set_cursor(*pygame.cursors.arrow)

    def waitRelease(self):
        while any(pygame.mouse.get_pressed()):
            pygame.event.pump()

    def checkCollision(self):
        mousePosition = pygame.mouse.get_pos()

        selected = None
        for i in range(len(self.rectangles)):
            if self.rectangles[i].collidepoint(mousePosition):
                selected = i

        return selected

class Text:

    def __init__(self, text, screen, font=Font().get(), colour=(0, 0, 0)):
        self.text = text
        self.screen = screen
        self.font = font
        self.colour = colour

    def write(self, position=(0, 0)):
        target = self.font.render(self.text, 1, self.colour)
        targetPosition = target.get_rect()
        targetPosition.centerx = position[0]
        targetPosition.centery = position[1]
        self.screen.blit(target, targetPosition)

        return target, targetPosition

    def wrapWrite(self, rect, lineSpace=1):
        """
        """
        space = self.font.render(" ", 1, self.colour)
        spaceSize = space.get_size()

        leftTop = [rect.left, rect.top]

        for sentence in self.text.split("\n"):
            for word in sentence.split(" "):
                word += " "

                target = self.font.render(word, 1, self.colour)
                targetSize = target.get_size()

                if (leftTop[0] + targetSize[0] > rect.right):
                    # new line
                    leftTop[0] = rect.left
                    leftTop[1] += spaceSize[1] * lineSpace

                self.screen.blit(target, leftTop)

                leftTop[0] += targetSize[0]

            leftTop[0] = rect.left
            leftTop[1] += spaceSize[1] * lineSpace
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
import pytest
import logging
import datetime
from pprint import pprint as pp

from scout.exceptions import (IntegrityError)

logger = logging.getLogger(__name__)

def test_add_cases(panel_database, case_obj):
    adapter = panel_database
    ## GIVEN an empty database (no cases)
    assert adapter.case_collection.find().count() == 0

    ## WHEN adding a new case to the database
    adapter.add_case(case_obj)

    ## THEN it should be populated with the new case
    result = adapter.cases()
    assert result.count() == 1
    for case in result:
        assert case['owner'] == case_obj['owner']

    logger.info("All cases checked")

def test_add_existing_case(panel_database, case_obj):
    adapter = panel_database
    ## GIVEN an empty database (no cases)
    assert adapter.cases().count() == 0

    adapter.add_case(case_obj)
    ## WHEN adding a existing case to the database
    with pytest.raises(IntegrityError):
    ## THEN it should raise integrity error
        adapter.add_case(case_obj)

def test_get_case(panel_database, case_obj):
    adapter = panel_database
    ## GIVEN an empty database (no cases)
    assert adapter.cases().count() == 0
    adapter.case_collection.insert_one(case_obj)
    logger.info("Testing to get case")

    ## WHEN retreiving an existing case from the database
    result = adapter.case(case_id=case_obj['_id'])
    ## THEN we should get the correct case
    assert result['owner'] == case_obj['owner']

def test_get_cases(real_panel_database, case_obj):
    adapter = real_panel_database
    ## GIVEN an empty database (no cases)
    assert adapter.cases().count() == 0
    adapter.case_collection.insert_one(case_obj)
    ## WHEN retreiving an existing case from the database
    result = adapter.cases()
    ## THEN we should get the correct case
    assert result.count() == 1

def test_get_cases_no_assignees(real_panel_database, case_obj):
    adapter = real_panel_database
    ## GIVEN an empty database (no cases)
    assert adapter.cases().count() == 0
    adapter.case_collection.insert_one(case_obj)
    ## WHEN retreiving an existing case from the database
    result = adapter.cases(name_query='john')
    ## THEN we should get the correct case
    assert result.count() == 0

def test_get_cases_display_name(real_panel_database, case_obj):
    adapter = real_panel_database
    ## GIVEN an empty database (no cases)
    assert adapter.cases().count() == 0
    adapter.case_collection.insert_one(case_obj)
    ## WHEN retreiving cases by partial display name
    result = adapter.cases(name_query='643')
    ## THEN we should get the correct case
    assert result.count() == 1

def test_get_cases_existing_individual(real_panel_database, case_obj):
    adapter = real_panel_database
    ## GIVEN an empty database (no cases)
    assert adapter.cases().count() == 0
    adapter.case_collection.insert_one(case_obj)
    ## WHEN retreiving cases by partial individual name
    result = adapter.cases(name_query='NA1288')
    ## THEN we should get the correct case
    assert result.count() == 1

def test_get_cases_assignees(real_panel_database, case_obj):
    adapter = real_panel_database
    ## GIVEN an empty database (no cases)
    assert adapter.cases().count() == 0

    user_obj = adapter.user_collection.find_one()
    case_obj['assignees'] = [user_obj['email']]
    adapter.case_collection.insert_one(case_obj)

    ## WHEN retreiving cases by partial individual name
    result = adapter.cases(name_query='john')
    ## THEN we should get the correct case
    assert result.count() == 1

def test_get_cases_non_existing_assignee(real_panel_database, case_obj):
    adapter = real_panel_database
    ## GIVEN an empty database (no cases)
    assert adapter.cases().count() == 0

    user_obj = adapter.user_collection.find_one()
    case_obj['assignees'] = [user_obj['email']]
    adapter.case_collection.insert_one(case_obj)

    ## WHEN retreiving cases by partial individual name
    result = adapter.cases(name_query='damien')
    ## THEN we should get the correct case
    assert result.count() == 0

def test_get_cases_non_existing_individual(real_panel_database, case_obj):
    adapter = real_panel_database
    ## GIVEN an empty database (no cases)
    assert adapter.cases().count() == 0
    adapter.case_collection.insert_one(case_obj)
    ## WHEN retreiving cases by partial display name
    result = adapter.cases(name_query='hello')
    ## THEN we should get the correct case
    assert result.count() == 0

def test_get_non_existing_case(panel_database, case_obj):
    adapter = panel_database
    # GIVEN an empty database (no cases)
    assert adapter.cases().count() == 0
    adapter.add_case(case_obj)
    logger.info("Testing to get case")

    # WHEN retreiving an non existing case from the database
    result = adapter.case(case_id='hello')
    # THEN we should get None back
    assert result is None

def test_delete_case(panel_database, case_obj):
    adapter = panel_database
    ## GIVEN an empty database (no cases)
    assert adapter.cases().count() == 0
    adapter.case_collection.insert_one(case_obj)
    assert adapter.cases().count() == 1
    logger.info("Testing to delete case")

    ## WHEN deleting a case from the database
    result = adapter.delete_case(case_id=case_obj['_id'])
    ## THEN there should be no cases left in the database
    assert adapter.cases().count() == 0

def test_update_case_collaborators(panel_database, case_obj):
    adapter = panel_database
    ## GIVEN an empty database (no cases)
    assert adapter.cases().count() == 0
    adapter.case_collection.insert_one(case_obj)
    assert adapter.cases().count() == 1
    assert len(adapter.case(case_obj['_id'])['collaborators']) == 1
    logger.info("Testing to update case")

    coll_1 = case_obj['collaborators'][0]
    coll_2 = 'test'
    coll_3 = 'test2'
    case_obj['collaborators'].append(coll_2)
    case_obj['collaborators'].append(coll_3)

    ## WHEN updating a case with new collaborators
    res = adapter.update_case(case_obj)

    ## THEN assert collaborator has been added
    assert len(res['collaborators']) == 3
    ## THEN assert all collaborators where added
    assert set(res['collaborators']) == set([coll_1, coll_2, coll_3])

def test_update_case_individuals(panel_database, case_obj):
    adapter = panel_database
    ## GIVEN an empty database (no cases)
    assert adapter.cases().count() == 0
    adapter.case_collection.insert_one(case_obj)
    assert adapter.cases().count() == 1
    logger.info("Testing to update case")

    res = adapter.case(case_obj['_id'])
    assert len(res['individuals']) == 3

    new_individuals = [{
        'individual_id': 'test',
        'display_name': 'test_name',
    }]
    case_obj['individuals'] = new_individuals
    ## WHEN updating a case with new individuals
    res = adapter.update_case(case_obj)
    ## THEN assert that 'individuals' has changed

    assert len(res['individuals']) == 1

def test_update_case_analysis_dates(panel_database, case_obj):
    adapter = panel_database
    ## GIVEN an empty database (no cases)
    assert adapter.cases().count() == 0
    adapter.case_collection.insert_one(case_obj)
    assert adapter.cases().count() == 1
    logger.info("Testing to update case")

    res = adapter.case(case_obj['_id'])
    old_date = res['analysis_date']
    old_update = res['updated_at']
    assert len(res['analysis_dates']) == 1

    new_analysis = datetime.datetime.now()
    case_obj['analysis_date'] = new_analysis
    ## WHEN updating a case with new analysis date
    res = adapter.update_case(case_obj)

    ## THEN assert that the new analysis date is added to 'analysis_dates'
    assert len(res['analysis_dates']) == 2
    ## THEN assert that 'analysis_date' is set to the new date
    assert res['analysis_date'] > old_date
    ## THEN assert that 'updated_at' has been changed
    assert res['updated_at'] > old_update

def test_update_case_rerun_status(panel_database, case_obj):
    adapter = panel_database
    case_obj['rerun_requested'] = True
    ## GIVEN an empty database (no cases)
    assert adapter.cases().count() == 0
    adapter.case_collection.insert_one(case_obj)
    assert adapter.cases().count() == 1
    logger.info("Testing to update case")

    res = adapter.case(case_obj['_id'])
    assert res['rerun_requested'] is True

    ## WHEN updating a case
    res = adapter.update_case(case_obj)

    ## THEN assert that 'rerun_requested' is set to False
    assert res['rerun_requested'] is False
<commit_msg>Updates test_case_handling<commit_after># -*- coding: utf-8 -*-
import pytest
import logging
import datetime
from pprint import pprint as pp

from scout.exceptions import (IntegrityError)

logger = logging.getLogger(__name__)

def test_add_cases(panel_database, case_obj):
    adapter = panel_database
    ## GIVEN an empty database (no cases)
    assert adapter.case_collection.find().count() == 0

    ## WHEN adding a new case to the database
    adapter._add_case(case_obj)

    ## THEN it should be populated with the new case
    result = adapter.cases()
    assert result.count() == 1
    for case in result:
        assert case['owner'] == case_obj['owner']

    logger.info("All cases checked")

def test_add_existing_case(panel_database, case_obj):
    adapter = panel_database
    ## GIVEN an empty database (no cases)
    assert adapter.cases().count() == 0

    adapter._add_case(case_obj)
    ## WHEN adding a existing case to the database
    with pytest.raises(IntegrityError):
    ## THEN it should raise integrity error
        adapter._add_case(case_obj)

def test_get_case(panel_database, case_obj):
    adapter = panel_database
    ## GIVEN an empty database (no cases)
    assert adapter.cases().count() == 0
    adapter.case_collection.insert_one(case_obj)
    logger.info("Testing to get case")

    ## WHEN retreiving an existing case from the database
    result = adapter.case(case_id=case_obj['_id'])
    ## THEN we should get the correct case
    assert result['owner'] == case_obj['owner']

def test_get_cases(real_panel_database, case_obj):
    adapter = real_panel_database
    ## GIVEN an empty database (no cases)
    assert adapter.cases().count() == 0
    adapter.case_collection.insert_one(case_obj)
    ## WHEN retreiving an existing case from the database
    result = adapter.cases()
    ## THEN we should get the correct case
    assert result.count() == 1

def test_get_cases_no_assignees(real_panel_database, case_obj):
    adapter = real_panel_database
    ## GIVEN an empty database (no cases)
    assert adapter.cases().count() == 0
    adapter.case_collection.insert_one(case_obj)
    ## WHEN retreiving an existing case from the database
    result = adapter.cases(name_query='john')
    ## THEN we should get the correct case
    assert result.count() == 0

def test_get_cases_display_name(real_panel_database, case_obj):
    adapter = real_panel_database
    ## GIVEN an empty database (no cases)
    assert adapter.cases().count() == 0
    adapter.case_collection.insert_one(case_obj)
    ## WHEN retreiving cases by partial display name
    result = adapter.cases(name_query='643')
    ## THEN we should get the correct case
    assert result.count() == 1

def test_get_cases_existing_individual(real_panel_database, case_obj):
    adapter = real_panel_database
    ## GIVEN an empty database (no cases)
    assert adapter.cases().count() == 0
    adapter.case_collection.insert_one(case_obj)
    ## WHEN retreiving cases by partial individual name
    result = adapter.cases(name_query='NA1288')
    ## THEN we should get the correct case
    assert result.count() == 1

def test_get_cases_assignees(real_panel_database, case_obj):
    adapter = real_panel_database
    ## GIVEN an empty database (no cases)
    assert adapter.cases().count() == 0

    user_obj = adapter.user_collection.find_one()
    case_obj['assignees'] = [user_obj['email']]
    adapter.case_collection.insert_one(case_obj)

    ## WHEN retreiving cases by partial individual name
    result = adapter.cases(name_query='john')
    ## THEN we should get the correct case
    assert result.count() == 1

def test_get_cases_non_existing_assignee(real_panel_database, case_obj):
    adapter = real_panel_database
    ## GIVEN an empty database (no cases)
    assert adapter.cases().count() == 0

    user_obj = adapter.user_collection.find_one()
    case_obj['assignees'] = [user_obj['email']]
    adapter.case_collection.insert_one(case_obj)

    ## WHEN retreiving cases by partial individual name
    result = adapter.cases(name_query='damien')
    ## THEN we should get the correct case
    assert result.count() == 0

def test_get_cases_non_existing_individual(real_panel_database, case_obj):
    adapter = real_panel_database
    ## GIVEN an empty database (no cases)
    assert adapter.cases().count() == 0
    adapter.case_collection.insert_one(case_obj)
    ## WHEN retreiving cases by partial display name
    result = adapter.cases(name_query='hello')
    ## THEN we should get the correct case
    assert result.count() == 0

def test_get_non_existing_case(panel_database, case_obj):
    adapter = panel_database
    # GIVEN an empty database (no cases)
    assert adapter.cases().count() == 0
    adapter._add_case(case_obj)
    logger.info("Testing to get case")

    # WHEN retreiving an non existing case from the database
    result = adapter.case(case_id='hello')
    # THEN we should get None back
    assert result is None

def test_delete_case(panel_database, case_obj):
    adapter = panel_database
    ## GIVEN an empty database (no cases)
    assert adapter.cases().count() == 0
    adapter.case_collection.insert_one(case_obj)
    assert adapter.cases().count() == 1
    logger.info("Testing to delete case")

    ## WHEN deleting a case from the database
    result = adapter.delete_case(case_id=case_obj['_id'])
    ## THEN there should be no cases left in the database
    assert adapter.cases().count() == 0

def test_update_case_collaborators(panel_database, case_obj):
    adapter = panel_database
    ## GIVEN an empty database (no cases)
    assert adapter.cases().count() == 0
    adapter.case_collection.insert_one(case_obj)
    assert adapter.cases().count() == 1
    assert len(adapter.case(case_obj['_id'])['collaborators']) == 1
    logger.info("Testing to update case")

    coll_1 = case_obj['collaborators'][0]
    coll_2 = 'test'
    coll_3 = 'test2'
    case_obj['collaborators'].append(coll_2)
    case_obj['collaborators'].append(coll_3)

    ## WHEN updating a case with new collaborators
    res = adapter.update_case(case_obj)

    ## THEN assert collaborator has been added
    assert len(res['collaborators']) == 3
    ## THEN assert all collaborators where added
    assert set(res['collaborators']) == set([coll_1, coll_2, coll_3])

def test_update_case_individuals(panel_database, case_obj):
    adapter = panel_database
    ## GIVEN an empty database (no cases)
    assert adapter.cases().count() == 0
    adapter.case_collection.insert_one(case_obj)
    assert adapter.cases().count() == 1
    logger.info("Testing to update case")

    res = adapter.case(case_obj['_id'])
    assert len(res['individuals']) == 3

    new_individuals = [{
        'individual_id': 'test',
        'display_name': 'test_name',
    }]
    case_obj['individuals'] = new_individuals
    ## WHEN updating a case with new individuals
    res = adapter.update_case(case_obj)
    ## THEN assert that 'individuals' has changed

    assert len(res['individuals']) == 1

def test_update_case_analysis_dates(panel_database, case_obj):
    adapter = panel_database
    ## GIVEN an empty database (no cases)
    assert adapter.cases().count() == 0
    adapter.case_collection.insert_one(case_obj)
    assert adapter.cases().count() == 1
    logger.info("Testing to update case")

    res = adapter.case(case_obj['_id'])
    old_date = res['analysis_date']
    old_update = res['updated_at']
    assert len(res['analysis_dates']) == 1

    new_analysis = datetime.datetime.now()
    case_obj['analysis_date'] = new_analysis
    ## WHEN updating a case with new analysis date
    res = adapter.update_case(case_obj)

    ## THEN assert that the new analysis date is added to 'analysis_dates'
    assert len(res['analysis_dates']) == 2
    ## THEN assert that 'analysis_date' is set to the new date
    assert res['analysis_date'] > old_date
    ## THEN assert that 'updated_at' has been changed
    assert res['updated_at'] > old_update

def test_update_case_rerun_status(panel_database, case_obj):
    adapter = panel_database
    case_obj['rerun_requested'] = True
    ## GIVEN an empty database (no cases)
    assert adapter.cases().count() == 0
    adapter.case_collection.insert_one(case_obj)
    assert adapter.cases().count() == 1
    logger.info("Testing to update case")

    res = adapter.case(case_obj['_id'])
    assert res['rerun_requested'] is True

    ## WHEN updating a case
    res = adapter.update_case(case_obj)

    ## THEN assert that 'rerun_requested' is set to False
    assert res['rerun_requested'] is False
<|endoftext|>


<commit_before>from autocontroller import AutoController

###### STATE MACHINE FOR NAVIGATIONAL STATES ########

class State:
    def run(self):
        assert 0, "run not implemented"
    def sensor_data_received(self, data, new_ir_right, new_ir_left):
        assert 0, "sensor_data_received not implemented"

class auto_control(State):

	auto_controller = AutoController()

	def sensor_data_received(self, data, new_ir_right, new_ir_left):
		##NOTE TO SELF: Kan inte reglera pa vanster sida nu!!!
		right_speed, left_speed = auto_control.auto_controller.auto_control(new_ir_right, new_ir_left, 'right')
		#Duration set to something quite high to mimic running forever (until next update)
		data['driver'].drive(left_speed, right_speed, 500)

	def run(self, data):
		#If sensor readings jump more than 5 mm we've discovered a turn
		print('distance diff: ' + str(data['ir_right'] - data['old_ir_right']))
		DISCONTINUITY_DIST = 20.0 #mm
		FACING_WALL_DIST = 30 #mm
		print('laser distance: ' + str(data['laser'].read_data()))
		if data['ir_right'] - data['old_ir_right'] >= DISCONTINUITY_DIST:
			print('changing to preparing for turn')
			data['driver'].prepare_for_turn()
			return before_turn()
		elif data['laser'].read_data() <=  FACING_WALL_DIST:
			data['driver'].turn_left()
			return turn()
		else:
			return auto_control()

class warmup(State):

	def sensor_data_received(self, data, new_ir_right, new_ir_left):
		return #Do nothing

	def run(self, data):
		if not data['driver'].driving():
			return auto_control()
		else:
			return warmup()

class before_turn(State):
	def sensor_data_received(self, data, new_ir_right, new_ir_left):
		return #Do nothing. Only auto control uses it

	def run(self, data):
		print('running before turn')
		if not data['driver'].driving():
			print('changing to turn')
			data['driver'].turn_right()
			return turn()
		else:
			return before_turn()

class after_turn(State):
	def sensor_data_received(self, data, new_ir_right, new_ir_left):
		return #Do nothing. Only auto control uses it

	def run(self, data):
		print('running after turn')
		if not data['driver'].driving():
			print('changing to auto control')
			return auto_control()
		else:
			return after_turn()

class turn(State):
	def sensor_data_received(self, data, new_ir_right, new_ir_left):
		return #Do nothing. Only auto control uses it

	def run(self, data):
		if not data['driver'].driving():
			print('changing to after turn')
			data['driver'].prepare_for_turn()
			return after_turn()
		else:
			return turn()

###### NAVIGATOR CLASS #######
class Navigator:

	def __init__(self, driver, laser):
		self.data = {'ir_left': 0,
					'ir_right': 0,
					'old_ir_right': 0,
					'old_ir_left': 0,
					'driver' : driver,
					'laser' : laser
					}

		#Stand still for 100 ms, waiting for sensors
		self.data['driver'].drive(0, 0, 1000)
		self.state = warmup()

	def sensor_data_received(self, new_ir_right, new_ir_left):
		self.state.sensor_data_received(self.data, new_ir_left, new_ir_right)
		self.data['old_ir_left'] = self.data['ir_left']
		self.data['old_ir_right'] = self.data['ir_right']
		self.data['ir_left'] = new_ir_left
		self.data['ir_right'] = new_ir_right

	#Runs the state. The states run method returns the next state
	def navigate(self):
		print('navigate')
		next_state = self.state.run(self.data)
		print(next_state)
		self.state = next_state<commit_msg>commented out laser<commit_after>from autocontroller import AutoController

###### STATE MACHINE FOR NAVIGATIONAL STATES ########

class State:
    def run(self):
        assert 0, "run not implemented"
    def sensor_data_received(self, data, new_ir_right, new_ir_left):
        assert 0, "sensor_data_received not implemented"

class auto_control(State):

	auto_controller = AutoController()

	def sensor_data_received(self, data, new_ir_right, new_ir_left):
		##NOTE TO SELF: Kan inte reglera pa vanster sida nu!!!
		right_speed, left_speed = auto_control.auto_controller.auto_control(new_ir_right, new_ir_left, 'right')
		#Duration set to something quite high to mimic running forever (until next update)
		data['driver'].drive(left_speed, right_speed, 500)

	def run(self, data):
		#If sensor readings jump more than 5 mm we've discovered a turn
		print('distance diff: ' + str(data['ir_right'] - data['old_ir_right']))
		print('laser distance: ' + str(data['laser'].read_data()))
		DISCONTINUITY_DIST = 20.0 #mm
		FACING_WALL_DIST = 30 #mm
		if data['ir_right'] - data['old_ir_right'] >= DISCONTINUITY_DIST:
			print('changing to preparing for turn')
			data['driver'].prepare_for_turn()
			return before_turn()
		#elif data['laser'].read_data() <=  FACING_WALL_DIST:
		#	data['driver'].turn_left()
		#	return turn()
		else:
			return auto_control()

class warmup(State):

	def sensor_data_received(self, data, new_ir_right, new_ir_left):
		return #Do nothing

	def run(self, data):
		if not data['driver'].driving():
			return auto_control()
		else:
			return warmup()

class before_turn(State):
	def sensor_data_received(self, data, new_ir_right, new_ir_left):
		return #Do nothing. Only auto control uses it

	def run(self, data):
		print('running before turn')
		if not data['driver'].driving():
			print('changing to turn')
			data['driver'].turn_right()
			return turn()
		else:
			return before_turn()

class after_turn(State):
	def sensor_data_received(self, data, new_ir_right, new_ir_left):
		return #Do nothing. Only auto control uses it

	def run(self, data):
		print('running after turn')
		if not data['driver'].driving():
			print('changing to auto control')
			return auto_control()
		else:
			return after_turn()

class turn(State):
	def sensor_data_received(self, data, new_ir_right, new_ir_left):
		return #Do nothing. Only auto control uses it

	def run(self, data):
		if not data['driver'].driving():
			print('changing to after turn')
			data['driver'].prepare_for_turn()
			return after_turn()
		else:
			return turn()

###### NAVIGATOR CLASS #######
class Navigator:

	def __init__(self, driver, laser):
		self.data = {'ir_left': 0,
					'ir_right': 0,
					'old_ir_right': 0,
					'old_ir_left': 0,
					'driver' : driver,
					'laser' : laser
					}

		#Stand still for 100 ms, waiting for sensors
		self.data['driver'].drive(0, 0, 1000)
		self.state = warmup()

	def sensor_data_received(self, new_ir_right, new_ir_left):
		self.state.sensor_data_received(self.data, new_ir_left, new_ir_right)
		self.data['old_ir_left'] = self.data['ir_left']
		self.data['old_ir_right'] = self.data['ir_right']
		self.data['ir_left'] = new_ir_left
		self.data['ir_right'] = new_ir_right

	#Runs the state. The states run method returns the next state
	def navigate(self):
		next_state = self.state.run(self.data)
		self.state = next_state<|endoftext|>


<commit_before>import logging
import datetime

from discord import User
from discord.ext import commands
from discord.ext.commands import Bot
from discord.ext.commands import Context

logger = logging.getLogger(__name__)

# Time where poll can only be closed by owner or channel managers
DEFAULT_LOCK_TIME = datetime.timedelta(hours=6)

class PollState:
    """State of a single poll.

     Attributes
    -----------
    owner : User
        Owner of this poll.
    question : str
        The poll question.
    choices : List[str]
        List of choices.
    voters : Set[User]
        Voters who have already voted in this poll.
    original_message: Message
        The message sent by the bot when the poll was started showing the state
        of the poll. It will be updated continuously as votes are submitted.
    """

    def __init__(self, owner: User, question: str, choices):
        self.owner = owner
        self.question = question
        self.choices = [{'name': choice, 'count': 0} for choice in choices]
        self.voters = set()
        self.original_message = None
        self.start_time = datetime.datetime.now()

    def __str__(self):
        choice_list = [
            '{}. {} - {} vote(s)'.format(i + 1, ans['name'], ans['count'])
            for i, ans in enumerate(self.choices)
        ]

        return '**POLL**: {}\n\n{}'.format(
            self.question, '\n'.join(choice_list))

    def end_result_str(self):
        max_votes = max(self.choices, key=lambda ans: ans['count'])['count']
        choice_list = [
            '{0}{1}. {2} - {3} vote(s){0}'.format(
                '**' if ans['count'] == max_votes else '',
                i + 1, ans['name'], ans['count'])
            for i, ans in enumerate(self.choices)
        ]

        return '**POLL RESULTS**: {}\n\n{}'.format(
            self.question, '\n'.join(choice_list))

    def is_valid_choice(self, ans_num):
        return 1 <= ans_num <= len(self.choices)

    def has_already_voted(self, user):
        return user in self.voters

    def has_lock_expired(self):
        return datetime.datetime.now() - self.start_time > DEFAULT_LOCK_TIME

class Poll:
    """Commands to create and manage polls."""

    def __init__(self, bot):
        self.bot = bot
        self.current_polls = {}

    @commands.group(pass_context=True)
    async def poll(self, ctx: Context):
        """Group of commands for poll management"""
        if ctx.invoked_subcommand is None:
            await self.bot.say('Unknown poll command.')

    @poll.command(pass_context=True)
    async def start(self, ctx: Context, question=None, *choices):
        """Start a poll.

        Syntax:
            /poll start "This is the question" "Choice 1" "Choice 2"
        """
        if question is None:
            await self.bot.reply('your poll needs a question.')
            return

        if len(choices) < 1:
            await self.bot.reply('your poll needs choices.')
            return

        channel = ctx.message.channel
        if channel in self.current_polls:
            await self.bot.reply('there is already a poll running in here.')
            return

        current_poll = PollState(ctx.message.author, question, choices)

        self.current_polls[channel] = current_poll

        message = await self.bot.say(current_poll)
        current_poll.original_message = message
        await self.bot.pin_message(message)

    @poll.command(pass_context=True, aliases=['end'])
    async def stop(self, ctx: Context):
        """Stop the current poll.

        When the poll is locked, only the poll owner or users who have the
        Manage Messages permission for the channel can stop the poll.
        """
        channel = ctx.message.channel
        if channel not in self.current_polls:
            await self.bot.reply('there is no poll currently running in here.')
            return

        current_poll = self.current_polls[channel]
        author = ctx.message.author
        if not (author == current_poll.owner or
                author.permissions_in(channel).manage_messages or
                current_poll.has_lock_expired()):
            await self.bot.reply("you can't end this poll yet.")
            return

        self.current_polls.pop(channel)
        await self.bot.unpin_message(current_poll.original_message)
        await self.bot.say(current_poll.end_result_str())

    @commands.command(pass_context=True)
    async def vote(self, ctx: Context, ans_num=None):
        """Vote for a choice. Must be a choice number."""
        channel = ctx.message.channel
        if channel not in self.current_polls:
            await self.bot.reply('there is no poll currently running in here.')
            return

        current_poll = self.current_polls[channel]

        if ans_num is None:
            await self.bot.reply('your vote must be between {}-{}.'.format(
                1, len(current_poll.choices)))
            return

        try:
            ans_num = int(ans_num)
        except ValueError:
            await self.bot.reply('your vote must be between {}-{}.'.format(
                1, len(current_poll.choices)))
            return

        author = ctx.message.author
        if current_poll.has_already_voted(author):
            await self.bot.reply('you have already voted in this poll.')
            return

        if not current_poll.is_valid_choice:
            await self.bot.reply('your vote must be between {}-{}.'.format(
                1, len(current_poll.choices)))

        current_poll.choices[ans_num - 1]['count'] += 1
        current_poll.voters.add(author)
        await self.bot.edit_message(current_poll.original_message,
                                    str(current_poll))
        await self.bot.reply('your vote has been counted.')

def setup(bot: Bot):
    bot.add_cog(Poll(bot))
<commit_msg>Combine identical replies<commit_after>import logging
import datetime

from discord import User
from discord.ext import commands
from discord.ext.commands import Bot
from discord.ext.commands import Context

logger = logging.getLogger(__name__)

# Time where poll can only be closed by owner or channel managers
DEFAULT_LOCK_TIME = datetime.timedelta(hours=6)

class PollState:
    """State of a single poll.

     Attributes
    -----------
    owner : User
        Owner of this poll.
    question : str
        The poll question.
    choices : List[str]
        List of choices.
    voters : Set[User]
        Voters who have already voted in this poll.
    original_message: Message
        The message sent by the bot when the poll was started showing the state
        of the poll. It will be updated continuously as votes are submitted.
    """

    def __init__(self, owner: User, question: str, choices):
        self.owner = owner
        self.question = question
        self.choices = [{'name': choice, 'count': 0} for choice in choices]
        self.voters = set()
        self.original_message = None
        self.start_time = datetime.datetime.now()

    def __str__(self):
        choice_list = [
            '{}. {} - {} vote(s)'.format(i + 1, ans['name'], ans['count'])
            for i, ans in enumerate(self.choices)
        ]

        return '**POLL**: {}\n\n{}'.format(
            self.question, '\n'.join(choice_list))

    def end_result_str(self):
        max_votes = max(self.choices, key=lambda ans: ans['count'])['count']
        choice_list = [
            '{0}{1}. {2} - {3} vote(s){0}'.format(
                '**' if ans['count'] == max_votes else '',
                i + 1, ans['name'], ans['count'])
            for i, ans in enumerate(self.choices)
        ]

        return '**POLL RESULTS**: {}\n\n{}'.format(
            self.question, '\n'.join(choice_list))

    def get_validated_choice(self, ans_num):
        try:
            ans_num = int(ans_num)
            return ans_num if 1 <= ans_num <= len(self.choices) else None
        except (TypeError, ValueError):
            return None

    def has_already_voted(self, user):
        return user in self.voters

    def has_lock_expired(self):
        return datetime.datetime.now() - self.start_time > DEFAULT_LOCK_TIME

class Poll:
    """Commands to create and manage polls."""

    def __init__(self, bot):
        self.bot = bot
        self.current_polls = {}

    @commands.group(pass_context=True)
    async def poll(self, ctx: Context):
        """Group of commands for poll management"""
        if ctx.invoked_subcommand is None:
            await self.bot.say('Unknown poll command.')

    @poll.command(pass_context=True)
    async def start(self, ctx: Context, question=None, *choices):
        """Start a poll.

        Syntax:
            /poll start "This is the question" "Choice 1" "Choice 2"
        """
        if question is None:
            await self.bot.reply('your poll needs a question.')
            return

        if len(choices) < 1:
            await self.bot.reply('your poll needs choices.')
            return

        channel = ctx.message.channel
        if channel in self.current_polls:
            await self.bot.reply('there is already a poll running in here.')
            return

        current_poll = PollState(ctx.message.author, question, choices)

        self.current_polls[channel] = current_poll

        message = await self.bot.say(current_poll)
        current_poll.original_message = message
        await self.bot.pin_message(message)

    @poll.command(pass_context=True, aliases=['end'])
    async def stop(self, ctx: Context):
        """Stop the current poll.

        When the poll is locked, only the poll owner or users who have the
        Manage Messages permission for the channel can stop the poll.
        """
        channel = ctx.message.channel
        if channel not in self.current_polls:
            await self.bot.reply('there is no poll currently running in here.')
            return

        current_poll = self.current_polls[channel]
        author = ctx.message.author
        if not (author == current_poll.owner or
                author.permissions_in(channel).manage_messages or
                current_poll.has_lock_expired()):
            await self.bot.reply("you can't end this poll yet.")
            return

        self.current_polls.pop(channel)
        await self.bot.unpin_message(current_poll.original_message)
        await self.bot.say(current_poll.end_result_str())

    @commands.command(pass_context=True)
    async def vote(self, ctx: Context, ans_num=None):
        """Vote for a choice. Must be a choice number."""
        channel = ctx.message.channel
        if channel not in self.current_polls:
            await self.bot.reply('there is no poll currently running in here.')
            return

        current_poll = self.current_polls[channel]

        ans_num = current_poll.get_validated_choice(ans_num)
        if ans_num is None:
            await self.bot.reply('your vote must be between {}-{}.'.format(
                1, len(current_poll.choices)))
            return

        author = ctx.message.author
        if current_poll.has_already_voted(author):
            await self.bot.reply('you have already voted in this poll.')
            return

        current_poll.choices[ans_num - 1]['count'] += 1
        current_poll.voters.add(author)
        await self.bot.edit_message(current_poll.original_message,
                                    str(current_poll))
        await self.bot.reply('your vote has been counted.')

def setup(bot: Bot):
    bot.add_cog(Poll(bot))
<|endoftext|>