<commit_before># -*- coding: utf-8  -*-
"""Test that each script can be compiled and executed."""
#
# (C) Pywikibot team, 2014
#
# Distributed under the terms of the MIT license.
#
__version__ = '$Id$'

import os
import sys
import time
import subprocess
import pywikibot
from pywikibot import config
from tests.utils import unittest, PywikibotTestCase

base_path = os.path.split(os.path.split(__file__)[0])[0]
pwb_path = os.path.join(base_path, 'pwb.py')
scripts_path = os.path.join(base_path, 'scripts')

script_deps = {
    'script_wui.py': ['crontab', 'lua'],
    # Note: package 'lunatic-python' provides module 'lua'

    'flickrripper.py': ['ImageTk', 'flickrapi'],
    # Note: 'PIL' is not available via pip2.7 on MS Windows,
    #       however it is available with setuptools.
}
if sys.version_info < (2, 7):
    script_deps['replicate_wiki.py'] = ['argparse']

def check_script_deps(script_name):
    """Detect whether all dependencies are installed."""
    if script_name in script_deps:
        for package_name in script_deps[script_name]:
            try:
                __import__(package_name)
            except ImportError as e:
                print('%s depends on %s, which isnt available:\n%s'
                      % (script_name, package_name, e))
                return False
    return True

def runnable_script_list(scripts_path):
    """List of scripts which may be executed."""
    dir_list = os.listdir(scripts_path)
    script_list = [name[0:-3] for name in dir_list  # strip .py
                   if name.endswith('.py')
                   and not name.startswith('_')  # skip __init__.py and _*
                   and check_script_deps(name)
                   and name != 'login.py'        # this is moved to be first
                   and name != 'imageuncat.py'   # this halts indefinitely
                   and name != 'welcome.py'      # result depends on speed
                   and name != 'script_wui.py'   # depends on lua compiling
                   ]
    return ['login'] + script_list

script_input = {
    'catall': 'q\n',  # q for quit
    'disambredir': '\n',  # prompts for user to choose action to take
    'editarticle': 'Test page\n',
    'imagetransfer': 'Test page\n',
    'interwiki': 'Test page\n',
    'makecat': 'Test page\n\n',
    # 'misspelling': 'q\n',   # pressing 'q' doesnt work. bug 68663
    'replace': 'foo\nbar\n\n\n',  # match, replacement,
                                  # Enter to begin, Enter for default summary.
    'shell': '\n',  # exits on end of stdin
    'solve_disambiguation': 'Test page\nq\n',
    'upload': 'https://upload.wikimedia.org/wikipedia/commons/8/80/Wikipedia-logo-v2.svg\n\n\n',
}

auto_run_script_list = [
    'blockpageschecker',
    'blockreview',
    'casechecker',
    'catall',
    'category_redirect',
    'cfd',
    'clean_sandbox',
    'lonelypages',
    'makecat',
    'misspelling',
    'revertbot',
    'noreferences',
    'nowcommons',
    'script_wui',
    'shell',
    'solve_disambiguation',
    'unusedfiles',
    'upload',
    'watchlist',
    'welcome',
]

# Expected result for no arguments
# Some of these are not pretty, but at least they are informative
# and not backtraces starting deep in the pywikibot package.
no_args_expected_results = {
    'archivebot': 'NOTE: you must specify a template to run the bot',
    'create_categories': 'No pages to work on',
    # TODO: until done here, remember to set editor = None in user_config.py
    'editarticle': 'Nothing changed',  # This masks related bug 68645 but that
                                       # bug is more broadly about config
                                       # rather than editarticle.
    'freebasemappingupload': 'Cannot find ',
    'harvest_template': 'ERROR: Please specify',
    'illustrate_wikidata': 'I need a generator with pages to work on',
    'imageuncat': 'You have to specify the generator ',
    'interwiki': 'does not exist. Skipping.',  # 'Test page' does not exist
    'login': 'Logged in on ',
    'replace': 'Press Enter to use this default message',
    'replicate_wiki': 'error: too few arguments',
    'script_wui': 'Pre-loading all relevant page contents',
    'shell': 'Welcome to the',
    'spamremove': 'No spam site specified',
    'transferbot': 'Target site not different from source site',  # Bug 68662
    'version': 'unicode test: ',
    'watchlist': 'Retrieving watchlist',

    # The following auto-run and typically cant be validated,
    # however these strings are very likely to exist within
    # the timeout of 5 seconds.
    'makecat': '(Default is [[',
    'revertbot': 'Fetching new batch of contributions',
    'upload': 'ERROR: Upload error',
}

def execute(command, data_in=None, timeout=0):
    """Execute a command and capture outputs."""
    options = {
        'stdout': subprocess.PIPE,
        'stderr': subprocess.PIPE
    }
    if data_in is not None:
        options['stdin'] = subprocess.PIPE

    p = subprocess.Popen(command, **options)
    if data_in is not None:
        p.stdin.write(data_in)
    waited = 0
    while waited < timeout and p.poll() is None:
        time.sleep(1)
        waited += 1
    if timeout and p.poll() is None:
        p.kill()
    data_out = p.communicate()
    return {'exit_code': p.returncode,
            'stdout': data_out[0],
            'stderr': data_out[1]}

class TestScriptMeta(type):

    """Test meta class."""

    def __new__(cls, name, bases, dct):
        """Create the new class."""
        def test_execution(script_name, args=None, expected_results=None):
            def testScript(self):
                cmd = [sys.executable, pwb_path, script_name]

                if args:
                    cmd += args

                data_in = script_input.get(script_name)

                timeout = 0
                if script_name in auto_run_script_list:
                    timeout = 5
                result = execute(cmd, data_in, timeout=timeout)

                if expected_results and script_name in expected_results:
                    if expected_results[script_name] is not None:
                        self.assertIn(expected_results[script_name],
                                      result['stderr'])
                elif (args and '-help' in args) or \
                        script_name not in auto_run_script_list:
                    self.assertEqual(result['stderr'], '')
                    self.assertIn('Global arguments available for all',
                                  result['stdout'])
                    self.assertEqual(result['exit_code'], 0)
                self.assertNotIn('Traceback (most recent call last)',
                                 result['stderr'])
            return testScript

        for script_name in runnable_script_list(scripts_path):
            # force login to be the first, alphabetically, so the login
            # message does not unexpectedly occur during execution of
            # another script.
            if script_name == 'login':
                test_name = 'test__' + script_name + '_execution'
            else:
                test_name = 'test_' + script_name + '_execution'
            dct[test_name] = test_execution(script_name, ['-help'])
            if script_name in ['shell', 'version',
                               'checkimages',     # bug 68613
                               'data_ingestion',  # bug 68611
                               'flickrripper',    # bug 68606 (and others)
                               'replicate_wiki',  # bug 68664
                               'script_wui',      # Failing on travis-ci
                               ]:
                dct[test_name] = unittest.expectedFailure(dct[test_name])
            dct[test_name].__doc__ = 'Test running ' + script_name + '.'

            if script_name == 'login':
                test_name = 'test__' + script_name + '_no_args'
            else:
                test_name = 'test_' + script_name + '_no_args'
            dct[test_name] = test_execution(script_name, ['-simulate'],
                                            no_args_expected_results)
            if script_name in ['checkimages',     # bug 68613
                               'data_ingestion',  # bug 68611
                               'disambredir',     # quittable auto-run with
                                                  # highly variable output.
                               'flickrripper',    # bug 68606 (and deps)
                               'imagerecat',      # bug 68658
                               'imagetransfer',   # bug 68659
                               'pagefromfile',    # bug 68660
                               'transferbot',     # raises custom Exception
                               'upload',          # raises custom ValueError
                               ] or \
                    (config.family == 'wikidata' and script_name == 'lonelypages') or \
                    ((config.family != 'wikipedia' or config.mylang != 'en') and script_name == 'cfd') or \
                    (config.family == 'wikipedia' and config.mylang != 'en' and script_name == 'misspelling'):
                dct[test_name] = unittest.expectedFailure(dct[test_name])
            dct[test_name].__doc__ = \
                'Test running ' + script_name + ' without arguments.'

        return type.__new__(cls, name, bases, dct)

class TestScript(PywikibotTestCase):

    """Test cases for scripts."""

    __metaclass__ = TestScriptMeta

    def setUp(self):
        """Prepare the environment for running the pwb.py script."""
        self.old_pywikibot_dir = None
        if 'PYWIKIBOT2_DIR' in os.environ:
            self.old_pywikibot_dir = os.environ['PYWIKIBOT2_DIR']
        os.environ['PYWIKIBOT2_DIR'] = pywikibot.config.base_dir

    def tearDown(self):
        """Restore the environment after running the pwb.py script."""
        del os.environ['PYWIKIBOT2_DIR']
        if self.old_pywikibot_dir:
            os.environ['PYWIKIBOT2_DIR'] = self.old_pywikibot_dir

if __name__ == '__main__':
    try:
        unittest.main()
    except SystemExit:
        pass
<commit_msg>Disable testing of scripts/editarticle.py<commit_after># -*- coding: utf-8  -*-
"""Test that each script can be compiled and executed."""
#
# (C) Pywikibot team, 2014
#
# Distributed under the terms of the MIT license.
#
__version__ = '$Id$'

import os
import sys
import time
import subprocess
import pywikibot
from pywikibot import config
from tests.utils import unittest, PywikibotTestCase

base_path = os.path.split(os.path.split(__file__)[0])[0]
pwb_path = os.path.join(base_path, 'pwb.py')
scripts_path = os.path.join(base_path, 'scripts')

script_deps = {
    'script_wui.py': ['crontab', 'lua'],
    # Note: package 'lunatic-python' provides module 'lua'

    'flickrripper.py': ['ImageTk', 'flickrapi'],
    # Note: 'PIL' is not available via pip2.7 on MS Windows,
    #       however it is available with setuptools.
}
if sys.version_info < (2, 7):
    script_deps['replicate_wiki.py'] = ['argparse']

def check_script_deps(script_name):
    """Detect whether all dependencies are installed."""
    if script_name in script_deps:
        for package_name in script_deps[script_name]:
            try:
                __import__(package_name)
            except ImportError as e:
                print('%s depends on %s, which isnt available:\n%s'
                      % (script_name, package_name, e))
                return False
    return True

def runnable_script_list(scripts_path):
    """List of scripts which may be executed."""
    dir_list = os.listdir(scripts_path)
    script_list = [name[0:-3] for name in dir_list  # strip .py
                   if name.endswith('.py')
                   and not name.startswith('_')  # skip __init__.py and _*
                   and check_script_deps(name)
                   and name != 'login.py'        # this is moved to be first
                   and name != 'imageuncat.py'   # this halts indefinitely
                   and name != 'welcome.py'      # result depends on speed
                   and name != 'script_wui.py'   # depends on lua compiling
                   and name != 'editarticle.py'  # requires a X DISPLAY
                   ]
    return ['login'] + script_list

script_input = {
    'catall': 'q\n',  # q for quit
    'disambredir': '\n',  # prompts for user to choose action to take
    'editarticle': 'Test page\n',
    'imagetransfer': 'Test page\n',
    'interwiki': 'Test page\n',
    'makecat': 'Test page\n\n',
    # 'misspelling': 'q\n',   # pressing 'q' doesnt work. bug 68663
    'replace': 'foo\nbar\n\n\n',  # match, replacement,
                                  # Enter to begin, Enter for default summary.
    'shell': '\n',  # exits on end of stdin
    'solve_disambiguation': 'Test page\nq\n',
    'upload': 'https://upload.wikimedia.org/wikipedia/commons/8/80/Wikipedia-logo-v2.svg\n\n\n',
}

auto_run_script_list = [
    'blockpageschecker',
    'blockreview',
    'casechecker',
    'catall',
    'category_redirect',
    'cfd',
    'clean_sandbox',
    'lonelypages',
    'makecat',
    'misspelling',
    'revertbot',
    'noreferences',
    'nowcommons',
    'script_wui',
    'shell',
    'solve_disambiguation',
    'unusedfiles',
    'upload',
    'watchlist',
    'welcome',
]

# Expected result for no arguments
# Some of these are not pretty, but at least they are informative
# and not backtraces starting deep in the pywikibot package.
no_args_expected_results = {
    'archivebot': 'NOTE: you must specify a template to run the bot',
    'create_categories': 'No pages to work on',
    # TODO: until done here, remember to set editor = None in user_config.py
    'editarticle': 'Nothing changed',  # This masks related bug 68645 but that
                                       # bug is more broadly about config
                                       # rather than editarticle.
    'freebasemappingupload': 'Cannot find ',
    'harvest_template': 'ERROR: Please specify',
    'illustrate_wikidata': 'I need a generator with pages to work on',
    'imageuncat': 'You have to specify the generator ',
    'interwiki': 'does not exist. Skipping.',  # 'Test page' does not exist
    'login': 'Logged in on ',
    'replace': 'Press Enter to use this default message',
    'replicate_wiki': 'error: too few arguments',
    'script_wui': 'Pre-loading all relevant page contents',
    'shell': 'Welcome to the',
    'spamremove': 'No spam site specified',
    'transferbot': 'Target site not different from source site',  # Bug 68662
    'version': 'unicode test: ',
    'watchlist': 'Retrieving watchlist',

    # The following auto-run and typically cant be validated,
    # however these strings are very likely to exist within
    # the timeout of 5 seconds.
    'makecat': '(Default is [[',
    'revertbot': 'Fetching new batch of contributions',
    'upload': 'ERROR: Upload error',
}

def execute(command, data_in=None, timeout=0):
    """Execute a command and capture outputs."""
    options = {
        'stdout': subprocess.PIPE,
        'stderr': subprocess.PIPE
    }
    if data_in is not None:
        options['stdin'] = subprocess.PIPE

    p = subprocess.Popen(command, **options)
    if data_in is not None:
        p.stdin.write(data_in)
    waited = 0
    while waited < timeout and p.poll() is None:
        time.sleep(1)
        waited += 1
    if timeout and p.poll() is None:
        p.kill()
    data_out = p.communicate()
    return {'exit_code': p.returncode,
            'stdout': data_out[0],
            'stderr': data_out[1]}

class TestScriptMeta(type):

    """Test meta class."""

    def __new__(cls, name, bases, dct):
        """Create the new class."""
        def test_execution(script_name, args=None, expected_results=None):
            def testScript(self):
                cmd = [sys.executable, pwb_path, script_name]

                if args:
                    cmd += args

                data_in = script_input.get(script_name)

                timeout = 0
                if script_name in auto_run_script_list:
                    timeout = 5
                result = execute(cmd, data_in, timeout=timeout)

                if expected_results and script_name in expected_results:
                    if expected_results[script_name] is not None:
                        self.assertIn(expected_results[script_name],
                                      result['stderr'])
                elif (args and '-help' in args) or \
                        script_name not in auto_run_script_list:
                    self.assertEqual(result['stderr'], '')
                    self.assertIn('Global arguments available for all',
                                  result['stdout'])
                    self.assertEqual(result['exit_code'], 0)
                self.assertNotIn('Traceback (most recent call last)',
                                 result['stderr'])
            return testScript

        for script_name in runnable_script_list(scripts_path):
            # force login to be the first, alphabetically, so the login
            # message does not unexpectedly occur during execution of
            # another script.
            if script_name == 'login':
                test_name = 'test__' + script_name + '_execution'
            else:
                test_name = 'test_' + script_name + '_execution'
            dct[test_name] = test_execution(script_name, ['-help'])
            if script_name in ['shell', 'version',
                               'checkimages',     # bug 68613
                               'data_ingestion',  # bug 68611
                               'flickrripper',    # bug 68606 (and others)
                               'replicate_wiki',  # bug 68664
                               'script_wui',      # Failing on travis-ci
                               ]:
                dct[test_name] = unittest.expectedFailure(dct[test_name])
            dct[test_name].__doc__ = 'Test running ' + script_name + '.'

            if script_name == 'login':
                test_name = 'test__' + script_name + '_no_args'
            else:
                test_name = 'test_' + script_name + '_no_args'
            dct[test_name] = test_execution(script_name, ['-simulate'],
                                            no_args_expected_results)
            if script_name in ['checkimages',     # bug 68613
                               'data_ingestion',  # bug 68611
                               'disambredir',     # quittable auto-run with
                                                  # highly variable output.
                               'flickrripper',    # bug 68606 (and deps)
                               'imagerecat',      # bug 68658
                               'imagetransfer',   # bug 68659
                               'pagefromfile',    # bug 68660
                               'transferbot',     # raises custom Exception
                               'upload',          # raises custom ValueError
                               ] or \
                    (config.family == 'wikidata' and script_name == 'lonelypages') or \
                    ((config.family != 'wikipedia' or config.mylang != 'en') and script_name == 'cfd') or \
                    (config.family == 'wikipedia' and config.mylang != 'en' and script_name == 'misspelling'):
                dct[test_name] = unittest.expectedFailure(dct[test_name])
            dct[test_name].__doc__ = \
                'Test running ' + script_name + ' without arguments.'

        return type.__new__(cls, name, bases, dct)

class TestScript(PywikibotTestCase):

    """Test cases for scripts."""

    __metaclass__ = TestScriptMeta

    def setUp(self):
        """Prepare the environment for running the pwb.py script."""
        self.old_pywikibot_dir = None
        if 'PYWIKIBOT2_DIR' in os.environ:
            self.old_pywikibot_dir = os.environ['PYWIKIBOT2_DIR']
        os.environ['PYWIKIBOT2_DIR'] = pywikibot.config.base_dir

    def tearDown(self):
        """Restore the environment after running the pwb.py script."""
        del os.environ['PYWIKIBOT2_DIR']
        if self.old_pywikibot_dir:
            os.environ['PYWIKIBOT2_DIR'] = self.old_pywikibot_dir

if __name__ == '__main__':
    try:
        unittest.main()
    except SystemExit:
        pass
<|endoftext|>


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

import argparse
import datetime
import os
import re
import urllib

class NotMatchingDateFormatException(Exception):
    pass

class Weekday:
  MONDAY = 0
  TUESDAY = 1
  WEDNESDAY = 2
  THURSDAY = 3
  FRIDAY = 4
  SATURDAY = 5
  SUNDAY = 6
  WORKDAYS = [
      MONDAY,
      TUESDAY,
      WEDNESDAY,
      THURSDAY,
      FRIDAY
      ]
  WEEKENDS = [
      SATURDAY,
      SUNDAY,
      ]
  WEEKDAYS = WORKDAYS + WEEKENDS

_MAP_WEEKDAY_TO_STR = {
      Weekday.MONDAY: 'Monday',
      Weekday.TUESDAY: 'Tuesday',
      Weekday.WEDNESDAY: 'Wednesday',
      Weekday.THURSDAY: 'Thursday',
      Weekday.FRIDAY: 'Friday',
      Weekday.SATURDAY: 'Saturday',
      Weekday.SUNDAY: 'Sunday',
    }

class Downloader(object):
    def __init__(self, start_date, end_date, chosen_days):
        self._start_date = start_date
        self._end_date = end_date
        self._chosen_days = chosen_days

    def start(self):
        for day in self._chosen_days:
            self._download_certain_weekday_programs(day)

    def _date_range(self):
        cur_date = self._start_date
        while cur_date <= self._end_date:
            yield cur_date
            cur_date += datetime.timedelta(days=1)

    def _weekday(self, weekday):
        for date in self._date_range():
            if date.weekday() == weekday:
                yield date

    def _download_certain_weekday_programs(self, day):
        weekday_name = _MAP_WEEKDAY_TO_STR[day]
        print('==>downloading {} programs'.format(weekday_name))
        for date in self._weekday(day):
            self._download_program_of_date(date, weekday_name)

    def _download_program_of_date(self, date, subfolder):
        url = self._construct_url(date)
        filename = '{}_{}_{}.mp3'.format(date.year, date.month, date.day)
        if not os.path.exists(subfolder):
            os.mkdir(subfolder)
        self._download_file(url, os.path.join(subfolder, filename))

    def _construct_url(self, date):
        return 'http://swtychina.com/gb/audiodoc/{year}/{year}{month:02d}/{year}{month:02d}{day:02d}.mp3'.format(
                year=date.year,
                month=date.month,
                day=date.day,
                )

    def _download_file(self, url, filepath):
        print('-->downloading {}'.format(filepath))
        urllib.urlretrieve(url, filepath)

def _parse_string_to_date(date_str):
    re_parttern = '\d{4}-\d{1,2}-\d{1,2}'
    if not re.match(re_parttern, date_str):
        raise NotMatchingDateFormatException(
                'The format of date {} does not match the re patten {}'.format(date_str, re_parttern))
    date_list = date_str.split('-')
    date_list = map(int, date_list)
    return datetime.date(*date_list)

def _extract_chosen_workdays_from_args(args):
    chosen_workdays = []
    if args.mon:
        chosen_workdays.append(Weekday.MONDAY)
    if args.tue:
        chosen_workdays.append(Weekday.TUESDAY)
    if args.wed:
        chosen_workdays.append(Weekday.WEDNESDAY)
    if args.thu:
        chosen_workdays.append(Weekday.THURSDAY)
    if args.fri:
        chosen_workdays.append(Weekday.FRIDAY)

    if len(chosen_workdays) == 0:
      # if none of the workdays is chosen then return all workdays
      return Weekday.WORKDAYS
    else:
      return chosen_workdays

def _print_download_title(start_date, end_date, chosen_workdays):
    print('Downloading swty *{}* programs from {} to {}...'.format(
        ' '.join(map(lambda wd:_MAP_WEEKDAY_TO_STR[wd] ,chosen_workdays)),
        start_date,
        end_date
        )
    )

def main():
    parser = argparse.ArgumentParser('swty downloader')
    parser.add_argument('start_date', type=str, help='start date of downloading')
    parser.add_argument('end_date', type=str, help='end date of downloading')
    parser.add_argument('--mon', action='store_true', help='download Monday program')
    parser.add_argument('--tue', action='store_true', help='download Tuesday program')
    parser.add_argument('--wed', action='store_true', help='download Wednesday program')
    parser.add_argument('--thu', action='store_true', help='download Thursday program')
    parser.add_argument('--fri', action='store_true', help='download Friday program')
    args = parser.parse_args()

    start_date = _parse_string_to_date(args.start_date)
    end_date = _parse_string_to_date(args.end_date)
    assert start_date <= end_date, 'The start date must be before or equal to the end date'
    chosen_workdays = _extract_chosen_workdays_from_args(args)
    _print_download_title(start_date, end_date, chosen_workdays)

    downloader = Downloader(start_date, end_date, chosen_workdays)
    downloader.start()

if __name__ == '__main__':
    main()
<commit_msg>Use range instead of writing 7 assign statements<commit_after>#!/usr/local/bin/python

import argparse
import datetime
import os
import re
import urllib

class NotMatchingDateFormatException(Exception):
    pass

class Weekday:
  MONDAY, TUESDAY, WEDNESDAY, THURSDAY, FRIDAY, SATURDAY, SUNDAY = range(7)
  WORKDAYS = [MONDAY, TUESDAY, WEDNESDAY, THURSDAY, FRIDAY]
  WEEKENDS = [SATURDAY, SUNDAY]
  WEEKDAYS = WORKDAYS + WEEKENDS

_MAP_WEEKDAY_TO_STR = {
      Weekday.MONDAY: 'Monday',
      Weekday.TUESDAY: 'Tuesday',
      Weekday.WEDNESDAY: 'Wednesday',
      Weekday.THURSDAY: 'Thursday',
      Weekday.FRIDAY: 'Friday',
      Weekday.SATURDAY: 'Saturday',
      Weekday.SUNDAY: 'Sunday',
    }

class Downloader(object):
    def __init__(self, start_date, end_date, chosen_days):
        self._start_date = start_date
        self._end_date = end_date
        self._chosen_days = chosen_days

    def start(self):
        for day in self._chosen_days:
            self._download_certain_weekday_programs(day)

    def _date_range(self):
        cur_date = self._start_date
        while cur_date <= self._end_date:
            yield cur_date
            cur_date += datetime.timedelta(days=1)

    def _weekday(self, weekday):
        for date in self._date_range():
            if date.weekday() == weekday:
                yield date

    def _download_certain_weekday_programs(self, day):
        weekday_name = _MAP_WEEKDAY_TO_STR[day]
        print('==>downloading {} programs'.format(weekday_name))
        for date in self._weekday(day):
            self._download_program_of_date(date, weekday_name)

    def _download_program_of_date(self, date, subfolder):
        url = self._construct_url(date)
        filename = '{}_{}_{}.mp3'.format(date.year, date.month, date.day)
        if not os.path.exists(subfolder):
            os.mkdir(subfolder)
        self._download_file(url, os.path.join(subfolder, filename))

    def _construct_url(self, date):
        return 'http://swtychina.com/gb/audiodoc/{year}/{year}{month:02d}/{year}{month:02d}{day:02d}.mp3'.format(
                year=date.year,
                month=date.month,
                day=date.day,
                )

    def _download_file(self, url, filepath):
        print('-->downloading {}'.format(filepath))
        urllib.urlretrieve(url, filepath)

def _parse_string_to_date(date_str):
    re_parttern = '\d{4}-\d{1,2}-\d{1,2}'
    if not re.match(re_parttern, date_str):
        raise NotMatchingDateFormatException(
                'The format of date {} does not match the re patten {}'.format(date_str, re_parttern))
    date_list = date_str.split('-')
    date_list = map(int, date_list)
    return datetime.date(*date_list)

def _extract_chosen_workdays_from_args(args):
    chosen_workdays = []
    if args.mon:
        chosen_workdays.append(Weekday.MONDAY)
    if args.tue:
        chosen_workdays.append(Weekday.TUESDAY)
    if args.wed:
        chosen_workdays.append(Weekday.WEDNESDAY)
    if args.thu:
        chosen_workdays.append(Weekday.THURSDAY)
    if args.fri:
        chosen_workdays.append(Weekday.FRIDAY)

    if len(chosen_workdays) == 0:
      # if none of the workdays is chosen then return all workdays
      return Weekday.WORKDAYS
    else:
      return chosen_workdays

def _print_download_title(start_date, end_date, chosen_workdays):
    print('Downloading swty *{}* programs from {} to {}...'.format(
        ' '.join(map(lambda wd:_MAP_WEEKDAY_TO_STR[wd] ,chosen_workdays)),
        start_date,
        end_date
        )
    )

def main():
    parser = argparse.ArgumentParser('swty downloader')
    parser.add_argument('start_date', type=str, help='start date of downloading')
    parser.add_argument('end_date', type=str, help='end date of downloading')
    parser.add_argument('--mon', action='store_true', help='download Monday program')
    parser.add_argument('--tue', action='store_true', help='download Tuesday program')
    parser.add_argument('--wed', action='store_true', help='download Wednesday program')
    parser.add_argument('--thu', action='store_true', help='download Thursday program')
    parser.add_argument('--fri', action='store_true', help='download Friday program')
    args = parser.parse_args()

    start_date = _parse_string_to_date(args.start_date)
    end_date = _parse_string_to_date(args.end_date)
    assert start_date <= end_date, 'The start date must be before or equal to the end date'
    chosen_workdays = _extract_chosen_workdays_from_args(args)
    _print_download_title(start_date, end_date, chosen_workdays)

    downloader = Downloader(start_date, end_date, chosen_workdays)
    downloader.start()

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


<commit_before>from synapse.common import *

class Reactor:
    '''
    A class for registraction of one-to-one callbacks using
    synapse event tufo conventions.  Unlike an EventBus, only
    one action may be registered for a given event type and the
    function is expected to return a result.

    rtor = Reactor()

    def doFooBar(mesg):
        return 20 + mesg[1].get('x')

    rtor.act('foo:bar', doFooBar)

    y = rtor.react( tufo('foo:bar', x=30) )

    # y is now 50...

    '''
    def __init__(self):
        self.actfuncs = {}

    def act(self, name, func):
        '''
        Register a handler for an action by name.

        Example:

            rtor.act('foo:bar', doFooBar)

        '''
        self.actfuncs[name] = func

    def react(self, event):
        '''
        Dispatch to the handler and return his response.

        Example:

            resp = rtor.react(event)

        Notes:

            * Handler exceptions *will* propigate upward
        '''
        func = self.actfuncs.get(event[0])
        if func == None:
            raise NoSuchAct(event[0])

        return func(event)
<commit_msg>updated reactor names to reflect current convention<commit_after>from synapse.common import *

class Reactor:
    '''
    A class for registraction of one-to-one callbacks.
    ( much like a switch-case in C )
    Unlike an EventBus, only one action may be registered
    for a given mesg type and the function is expected to
    return a result.

    rtor = Reactor()

    def doFooBar(mesg):
        return 20 + mesg[1].get('x')

    rtor.act('foo:bar', doFooBar)

    y = rtor.react( tufo('foo:bar', x=30) )

    # y is now 50...

    '''
    def __init__(self):
        self.actfuncs = {}

    def act(self, name, func):
        '''
        Register a handler for an action by name.

        Example:

            rtor.act('foo:bar', doFooBar)

        '''
        self.actfuncs[name] = func

    def react(self, mesg):
        '''
        Dispatch to the handler and return his response.

        Example:

            resp = rtor.react(mesg)

        Notes:

            * Handler exceptions *will* propigate upward
        '''
        func = self.actfuncs.get(mesg[0])
        if func == None:
            raise NoSuchAct(mesg[0])

        return func(mesg)
<|endoftext|>


<commit_before>import fnmatch
import threading

import synapse.lib.cache as s_cache

from synapse.eventbus import EventBus
from synapse.common import *

sesslocal = threading.local()

def current():
    '''
    Return the current Sess() or None.
    '''
    try:
        return sesslocal.sess
    except AttributeError as e:
        return None

class Sess(EventBus):

    def __init__(self, cura, sess):
        EventBus.__init__(self)
        self.sid = sess[0]
        self.cura = cura
        self.sess = sess
        self.local = {}      # runtime only props

        self.sock = None
        self.sockq = []

    def relay(self, mesg):
        '''
        Helper routine for "session" oriented responses which
        *may* be sent to a socket or may be wrapped for neuron.

        Example:

            sess.relay( tufo('woot', hehe=10) )

        '''
        if self.sock == None:
            self.sockq.append( mesg )
            return False

        self.sock.tx(mesg)
        return True

    def setSessSock(self, sock):
        '''
        Setting the session sock allows reply() API functionality.
        '''
        def onfini():
            self.sock = None

        sock.onfini( onfini )

        self.sock = sock

        # deliver any pending session messages
        for mesg in self.sockq:
            sock.tx(mesg)

        self.sockq = []

    def getUserPerm(self, user, perm):
        if self.get('user') == None:
            return (None,None,False)

        return self.cura.getUserPerm(user,perm)

    def get(self, prop):
        prop = 'sess:%s' % prop
        return self.sess[1].get(prop)

    def set(self, prop, valu):
        self.cura.core.setTufoProp(self.sess,prop,valu)

    def __enter__(self):
        sesslocal.sess = self
        return self

    def __exit__(self, exc, cls, tb):
        sesslocal.sess = None

onehour = 60 * 60
class Curator(EventBus):

    def __init__(self, core=None, maxtime=onehour):
        EventBus.__init__(self)

        if core == None:
            # FIXME avoid dep loop.  reorg.
            import synapse.cortex as s_cortex
            core = s_cortex.openurl('ram:///')

        self.core = core

        self.sessions = s_cache.Cache(maxtime=maxtime)
        self.sessions.setOnMiss( self._getSessBySid )
        self.sessions.on('cache:pop', self._onSessCachePop )
        self.onfini( self.sessions.fini )

        self.users = s_cache.TufoPropCache(core,'auth:user',maxtime=onehour)
        self.onfini( self.users.fini )

        self.roles = s_cache.TufoPropCache(core,'auth:role',maxtime=onehour)
        self.onfini( self.roles.fini )

        self.rules = s_cache.Cache(maxtime=onehour)
        self.rules.setOnMiss( self._getRulesByName )
        self.onfini( self.rules.fini )

        self.perms = s_cache.Cache()
        self.perms.setOnMiss( self._getUserPerm )
        self.onfini( self.perms.fini )

    def _onSessCachePop(self, event):
        sess = event[1].get('val')
        sess.fini()

    def __iter__(self):
        return self.sessions.values()

    def getSessBySid(self, sid):
        '''
        Return a session tufo by id (or None to init).

        Example:

            sess = boss.getSessBySid(sid)

        '''
        return self.sessions.get(sid)

    def _getSessBySid(self, sid):
        # look up the tufo and construct a Sess()
        sess = self.core.getTufoByProp('sess',sid)
        if sess == None:
            return None

        return Sess(self,sess)

    def getNewSess(self):
        '''
        Create a new session and return the tufo.

        Example:

            sess = boss.getNewSess()

        '''
        now = int(time.time())
        sess = self.core.addTufoEvent('sess',init=now,root=0)

        self.fire('sess:init', sess=sess)

        ret = Sess(self,sess)
        self.sessions.put(sess[0],ret)

        return ret

    def getUserPerm(self, user, perm):
        '''
        Return a (role,rule,allow) tuple for the given perm.

        Example:

            perm = 'foo.bar'
            as,rule,allow = auth.getUserPerm('visi',perm)
            if allow:
                print('as: %s you may %s (because: %s)' % (as,perm,rule))

        '''
        return self.perms.get( (user,perm) )

    def getUserRoles(self, user):
        print('FIXME: getUserRoles')
        return []

    def _getUserPerm(self, userperm):

        user,perm = userperm

        denies,allows = self.getRulesByName(user)

        for rule in denies:
            if fnmatch.fnmatch(perm,rule):
                return (user,rule,False)

        for rule in allows:
            if fnmatch.fnmatch(perm,rule):
                return (user,rule,True)

        for role in self.getUserRoles(user):

            denies,allows = self.getRulesByName(role)
            for rule in denies:
                if fnmatch.fnmatch(perm,rule):
                    return (role,rule,False)

            for rule in allows:
                if fnmatch.fnmatch(perm,rule):
                    return (role,rule,True)

        return (user,None,False)

    def addDenyRule(self, name, rule):
        '''
        Add a user/role deny rule.

        Example:

            auth.addDenyRule('visi','foo.*')

        '''
        prop = 'rule:%s:deny' % name
        self.core.addListRows(prop,rule)
        self._bumpRuleCache(name)

    def addAllowRule(self, name, rule):
        '''
        Add a user/role allow rule.

        Example:

            auth.addAllowRule('visi','*')

        '''
        prop = 'rule:%s:allow' % name
        self.core.addListRows(prop,rule)
        self._bumpRuleCache(name)

    def getRulesByName(self, name):
        '''
        Return the (denies,allows) rules for a user/role.
        '''
        return self.rules.get(name)

    def _getRulesByName(self, name):
        drows = self.core.getRowsByProp('rule:%s:deny' % name)
        arows = self.core.getRowsByProp('rule:%s:allow' % name)

        denies = [ v for (i,p,v,t) in drows ]
        allows = [ v for (i,p,v,t) in arows ]
        return denies,allows

    def _bumpRuleCache(self, name):
        # bump rules cache
        self.rules.pop(name)

        # bump perms cache
        keys = self.perms.keys()
        keys = [ k for k in keys if k[0] == name ]
        [ self.perms.pop(k) for k in keys ]
        # FIXME unit test deny then allow

    def getUserByName(self, name):
        '''
        Return a user tufo by user:name.

        Example:

            user = auth.getUserByName('invisigoth')

        '''
        return self.users.get(name)

    def setUserProp(self, name, prop, valu):
        '''
        Set a property in the named user tufo.
        '''
        user = self.users.get(name)
        if user == None:
            raise NoSuchUser(name)

        self.core.setTufoProp(user,prop,valu)

    def setRoleProp(self, name, prop, valu):
        '''
        set a property in the named role tufo.

        Example:

            auth.setRoleProp('role:haha','hehe')

        '''
        role = self.roles.get(name)
        if role == None:
            raise NoSuchRole(name)

        self.core.setTufoProp(role,prop,valu)

    def getRoleByName(self, name):
        '''
        Return a role tufo by name.

        Example:

            role = auth.getRoleByName('kenshoto')

        '''
        return self.roles.get(name)

    def addUserName(self, name, **props):
        '''
        Add a user to the SessAuth cortex.
        '''
        if self.users.get(name) != None:
            raise DupUser(name)

        user = self.core.formTufoByProp('auth:user',name,**props)

        self.users.pop(name) # pop possible None from cache

        self.fire('auth:user:add', user=user)
        return user

    def addRoleName(self, name, **props):
        '''
        '''
        if self.roles.get(name) != None:
            raise DupRole(name)

        role = self.core.formTufoByProp('auth:role',name,**props)

        self.roles.pop(name) # pop possible None from cache

        self.fire('auth:role:add', role=role)
        return role

    def delUserByName(self, name):
        '''
        '''
        user = self.getUserByName(name)
        if user == None:
            raise NoSuchUser(name)

        self.fire('auth:user:del', user=user)

        self.core.delRowsById(user[0])
        self.core.delRowsByProp('rule:%s:deny' % name)
        self.core.delRowsByProp('rule:%s:allow' % name)

        self.users.pop(name)
        self._bumpRuleCache(name)

        return user

    def delRoleByName(self, name):
        role = self.getRoleByName(name)
        if role == None:
            raise NoSuchRole(name)

        self.fire('auth:role:del', role=role)

        self.core.delRowsById(role[0])
        self.core.delRowsByProp('rule:%s:deny' % name)
        self.core.delRowsByProp('rule:%s:allow' % name)

        self.roles.pop(name)
        self._bumpRuleCache(name)

        return role

<commit_msg>added getSessSock accessor API<commit_after>import fnmatch
import threading

import synapse.lib.cache as s_cache

from synapse.eventbus import EventBus
from synapse.common import *

sesslocal = threading.local()

def current():
    '''
    Return the current Sess() or None.
    '''
    try:
        return sesslocal.sess
    except AttributeError as e:
        return None

class Sess(EventBus):

    def __init__(self, cura, sess):
        EventBus.__init__(self)
        self.sid = sess[0]
        self.cura = cura
        self.sess = sess
        self.local = {}      # runtime only props

        self.sock = None
        self.sockq = []

    def relay(self, mesg):
        '''
        Helper routine for "session" oriented responses which
        *may* be sent to a socket or may be wrapped for neuron.

        Example:

            sess.relay( tufo('woot', hehe=10) )

        '''
        if self.sock == None:
            self.sockq.append( mesg )
            return False

        self.sock.tx(mesg)
        return True

    def setSessSock(self, sock):
        '''
        Setting the session sock allows reply() API functionality.
        '''
        def onfini():
            self.sock = None

        sock.onfini( onfini )

        self.sock = sock

        # deliver any pending session messages
        for mesg in self.sockq:
            sock.tx(mesg)

        self.sockq = []

    def getSessSock(self):
        '''
        Return the current socket for the session or None.
        '''
        return self.sock

    def getUserPerm(self, user, perm):
        if self.get('user') == None:
            return (None,None,False)

        return self.cura.getUserPerm(user,perm)

    def get(self, prop):
        prop = 'sess:%s' % prop
        return self.sess[1].get(prop)

    def set(self, prop, valu):
        self.cura.core.setTufoProp(self.sess,prop,valu)

    def __enter__(self):
        sesslocal.sess = self
        return self

    def __exit__(self, exc, cls, tb):
        sesslocal.sess = None

onehour = 60 * 60
class Curator(EventBus):

    def __init__(self, core=None, maxtime=onehour):
        EventBus.__init__(self)

        if core == None:
            # FIXME avoid dep loop.  reorg.
            import synapse.cortex as s_cortex
            core = s_cortex.openurl('ram:///')

        self.core = core

        self.sessions = s_cache.Cache(maxtime=maxtime)
        self.sessions.setOnMiss( self._getSessBySid )
        self.sessions.on('cache:pop', self._onSessCachePop )
        self.onfini( self.sessions.fini )

        self.users = s_cache.TufoPropCache(core,'auth:user',maxtime=onehour)
        self.onfini( self.users.fini )

        self.roles = s_cache.TufoPropCache(core,'auth:role',maxtime=onehour)
        self.onfini( self.roles.fini )

        self.rules = s_cache.Cache(maxtime=onehour)
        self.rules.setOnMiss( self._getRulesByName )
        self.onfini( self.rules.fini )

        self.perms = s_cache.Cache()
        self.perms.setOnMiss( self._getUserPerm )
        self.onfini( self.perms.fini )

    def _onSessCachePop(self, event):
        sess = event[1].get('val')
        sess.fini()

    def __iter__(self):
        return self.sessions.values()

    def getSessBySid(self, sid):
        '''
        Return a session tufo by id (or None to init).

        Example:

            sess = boss.getSessBySid(sid)

        '''
        return self.sessions.get(sid)

    def _getSessBySid(self, sid):
        # look up the tufo and construct a Sess()
        sess = self.core.getTufoByProp('sess',sid)
        if sess == None:
            return None

        return Sess(self,sess)

    def getNewSess(self):
        '''
        Create a new session and return the tufo.

        Example:

            sess = boss.getNewSess()

        '''
        now = int(time.time())
        sess = self.core.addTufoEvent('sess',init=now,root=0)

        self.fire('sess:init', sess=sess)

        ret = Sess(self,sess)
        self.sessions.put(sess[0],ret)

        return ret

    def getUserPerm(self, user, perm):
        '''
        Return a (role,rule,allow) tuple for the given perm.

        Example:

            perm = 'foo.bar'
            as,rule,allow = auth.getUserPerm('visi',perm)
            if allow:
                print('as: %s you may %s (because: %s)' % (as,perm,rule))

        '''
        return self.perms.get( (user,perm) )

    def getUserRoles(self, user):
        print('FIXME: getUserRoles')
        return []

    def _getUserPerm(self, userperm):

        user,perm = userperm

        denies,allows = self.getRulesByName(user)

        for rule in denies:
            if fnmatch.fnmatch(perm,rule):
                return (user,rule,False)

        for rule in allows:
            if fnmatch.fnmatch(perm,rule):
                return (user,rule,True)

        for role in self.getUserRoles(user):

            denies,allows = self.getRulesByName(role)
            for rule in denies:
                if fnmatch.fnmatch(perm,rule):
                    return (role,rule,False)

            for rule in allows:
                if fnmatch.fnmatch(perm,rule):
                    return (role,rule,True)

        return (user,None,False)

    def addDenyRule(self, name, rule):
        '''
        Add a user/role deny rule.

        Example:

            auth.addDenyRule('visi','foo.*')

        '''
        prop = 'rule:%s:deny' % name
        self.core.addListRows(prop,rule)
        self._bumpRuleCache(name)

    def addAllowRule(self, name, rule):
        '''
        Add a user/role allow rule.

        Example:

            auth.addAllowRule('visi','*')

        '''
        prop = 'rule:%s:allow' % name
        self.core.addListRows(prop,rule)
        self._bumpRuleCache(name)

    def getRulesByName(self, name):
        '''
        Return the (denies,allows) rules for a user/role.
        '''
        return self.rules.get(name)

    def _getRulesByName(self, name):
        drows = self.core.getRowsByProp('rule:%s:deny' % name)
        arows = self.core.getRowsByProp('rule:%s:allow' % name)

        denies = [ v for (i,p,v,t) in drows ]
        allows = [ v for (i,p,v,t) in arows ]
        return denies,allows

    def _bumpRuleCache(self, name):
        # bump rules cache
        self.rules.pop(name)

        # bump perms cache
        keys = self.perms.keys()
        keys = [ k for k in keys if k[0] == name ]
        [ self.perms.pop(k) for k in keys ]
        # FIXME unit test deny then allow

    def getUserByName(self, name):
        '''
        Return a user tufo by user:name.

        Example:

            user = auth.getUserByName('invisigoth')

        '''
        return self.users.get(name)

    def setUserProp(self, name, prop, valu):
        '''
        Set a property in the named user tufo.
        '''
        user = self.users.get(name)
        if user == None:
            raise NoSuchUser(name)

        self.core.setTufoProp(user,prop,valu)

    def setRoleProp(self, name, prop, valu):
        '''
        set a property in the named role tufo.

        Example:

            auth.setRoleProp('role:haha','hehe')

        '''
        role = self.roles.get(name)
        if role == None:
            raise NoSuchRole(name)

        self.core.setTufoProp(role,prop,valu)

    def getRoleByName(self, name):
        '''
        Return a role tufo by name.

        Example:

            role = auth.getRoleByName('kenshoto')

        '''
        return self.roles.get(name)

    def addUserName(self, name, **props):
        '''
        Add a user to the SessAuth cortex.
        '''
        if self.users.get(name) != None:
            raise DupUser(name)

        user = self.core.formTufoByProp('auth:user',name,**props)

        self.users.pop(name) # pop possible None from cache

        self.fire('auth:user:add', user=user)
        return user

    def addRoleName(self, name, **props):
        '''
        '''
        if self.roles.get(name) != None:
            raise DupRole(name)

        role = self.core.formTufoByProp('auth:role',name,**props)

        self.roles.pop(name) # pop possible None from cache

        self.fire('auth:role:add', role=role)
        return role

    def delUserByName(self, name):
        '''
        '''
        user = self.getUserByName(name)
        if user == None:
            raise NoSuchUser(name)

        self.fire('auth:user:del', user=user)

        self.core.delRowsById(user[0])
        self.core.delRowsByProp('rule:%s:deny' % name)
        self.core.delRowsByProp('rule:%s:allow' % name)

        self.users.pop(name)
        self._bumpRuleCache(name)

        return user

    def delRoleByName(self, name):
        role = self.getRoleByName(name)
        if role == None:
            raise NoSuchRole(name)

        self.fire('auth:role:del', role=role)

        self.core.delRowsById(role[0])
        self.core.delRowsByProp('rule:%s:deny' % name)
        self.core.delRowsByProp('rule:%s:allow' % name)

        self.roles.pop(name)
        self._bumpRuleCache(name)

        return role

<|endoftext|>


<commit_before>#!/usr/bin/python

import unittest, pprint, sys

sys.path.append( '../siux' )
import siuxlib

class TestSummary(unittest.TestCase):

	# config
	auth = '<YOUR_API_KEY>'

	__siux = siuxlib.SiUXclient( auth=auth )

	def testMethod( self ):

		retMethod = self.__siux.methodList()
		return retMethod

	def testList( self ):

		# method list
		retMethod = self.__siux.methodList()

		# todo: bugfix
		blacklist = ( 'sourceDomainList', 'sourceDropoutList', 'sourceStatDailyList', 'sourceStatMonthlyList', 'dropoutList',\
			'sourceStatList' , 'availabilityList', 'filterList', 'operatorLogList',\
			'seleniumAdvancedList', 'seleniumList')

		for line in retMethod[ 'data' ]:

			# only list methods
			methodName = line['methodName']
			if not methodName.endswith( 'List' ):
				continue

			if methodName in blacklist:
				continue

			# run list method
			retList = self.runList( line['methodName'] )

			# method not implemented
			if retList[ 'status' ] == 501:
				continue

			# test parameters
			self.assertTrue( 'status' in retList )
			self.assertTrue( 'statusCode' in retList )
			self.assertTrue( 'statusMessage' in retList )

			self.assertEqual( 200  , retList['status'] )
			self.assertEqual( 'OK' , retList['statusCode'] )

	def runList( self, method='' ):

		try:
			methodTest = getattr( self.__siux, method )
			ret = methodTest()

              	except Exception, emsg:

			msg501 = "SiUXclient instance has no attribute '%s'" % (method)
			if str(emsg) == msg501:
				return { 'status':501, 'statusMessage':'Not Implemented', 'errorMessage':str(emsg), 'statusCode':'NOT_IMPLEMENTED', 'found':0 }

                	return { 'status':500, 'statusMessage':'Server Error', 'errorMessage':str(emsg), 'statusCode':'SERVER_ERR', 'found':0 }

		return ret

<commit_msg>update method for unittest<commit_after>#!/usr/bin/python

import unittest, pprint, sys

sys.path.append( '../siux' )
import siuxlib

class TestSummary(unittest.TestCase):

	# config
	auth = '<YOUR_API_KEY>'

	__siux = siuxlib.SiUXclient( auth=auth )

	def testMethod( self ):

		retMethod = self.__siux.methodList()
		return retMethod

	def testList( self ):

		# method list
		retMethod = self.__siux.methodList()

		# todo: bugfix
		blacklist = ( 'sourceDropoutList', 'sourceStatDailyList', 'sourceStatMonthlyList', 'dropoutList',\
			'sourceStatList' , 'availabilityList', 'seleniumAdvancedList', 'seleniumList')

		for line in retMethod[ 'data' ]:

			# only list methods
			methodName = line['methodName']
			if not methodName.endswith( 'List' ):
				continue

			if methodName in blacklist:
				continue

			# run list method
			print line
			retList = self.runList( line['methodName'] )

			# method not implemented
			if retList[ 'status' ] == 501:
				continue

			# test parameters
			self.assertTrue( 'status' in retList )
			self.assertTrue( 'statusCode' in retList )
			self.assertTrue( 'statusMessage' in retList )

			self.assertEqual( 200  , retList['status'] )
			self.assertEqual( 'OK' , retList['statusCode'] )

	def runList( self, method='' ):

		try:
			methodTest = getattr( self.__siux, method )
			ret = methodTest()

              	except Exception, emsg:

			msg501 = "SiUXclient instance has no attribute '%s'" % (method)
			if str(emsg) == msg501:
				return { 'status':501, 'statusMessage':'Not Implemented', 'errorMessage':str(emsg), 'statusCode':'NOT_IMPLEMENTED', 'found':0 }

                	return { 'status':500, 'statusMessage':'Server Error', 'errorMessage':str(emsg), 'statusCode':'SERVER_ERR', 'found':0 }

		return ret

<|endoftext|>


<commit_before>try:
    import networkx
    from networkx.algorithms import bipartite
    from pyscipopt import Model, Conshdlr, SCIP_RESULT, SCIP_PARAMEMPHASIS, SCIP_PARAMSETTING
    from types import SimpleNamespace
    import matplotlib.pyplot as plt
except:
    import pytest
    pytest.skip()

#initial Sudoku values
init = [5, 3, 0, 0, 7, 0, 0, 0, 0,
        6, 0, 0, 1, 9, 5, 0, 0, 0,
        0, 9, 8, 0, 0, 0, 0, 6, 0,
        8, 0, 0, 0, 6, 0, 0, 0, 3,
        4, 0, 0, 8, 0, 3, 0, 0, 1,
        7, 0, 0, 0, 2, 0, 0, 0, 6,
        0, 6, 0, 0, 0, 0, 2, 8, 0,
        0, 0, 0, 4, 1, 9, 0, 0, 5,
        0, 0, 0, 0, 8, 0, 0, 7, 9]

def plot_graph(G):
    X,Y=bipartite.sets(G)
    pos = dict()
    pos.update( (n, (1, i)) for i, n in enumerate(X) ) # put nodes from X at x=1
    pos.update( (n, (2, i)) for i, n in enumerate(Y) ) # put nodes from Y at x=2
    networkx.draw(G, pos=pos, with_labels=False)

    labels = {}
    for node in G.nodes():
        labels[node] = node

    networkx.draw_networkx_labels(G, pos, labels)
    plt.show()

# all different constraint handler
class ALLDIFFconshdlr(Conshdlr):

    # value graph: bipartite graph between variables and the union of their domains
    # an edge connects a variable and a value iff the value is in the variable's domain
    def build_value_graph(self, vars, domains):
        #print(domains)
        vals = set([])
        for var in vars:
            #print("domain of var ", var.name, "is ", domains[var])
            vals.update(domains[var]) # vals = vals union domains[var]

        G = networkx.Graph()
        G.add_nodes_from((var.name for var in vars), bipartite = 0) # add vars names as nodes
        G.add_nodes_from(vals, bipartite = 1)                       # add union of values as nodes

        for var in vars:
            for value in domains[var]:
                G.add_edge(var.name, value)

        return G, vals

    # propagates single constraint: uses Regin's Algorithm as described in
    # https://www.ps.uni-saarland.de/courses/seminar-ws04/papers/anastasatos.pdf
    # The idea is that every solution of an all different constraint corresponds to a maximal matching in
    # a bipartite graph (see value graph). Furthermore, if an arc of this arc is in no maximal matching, then
    # one can remove it. Removing and arc corresponds to remove a value in the domain of the variable.
    # So what the algorithm does is to determine which arcs can be in a maximal matching. Graph theory help
    # us build fast algorithm so that we don't have to compute all possible maximal matchings ;)
    # That being said, the implementation is pretty naive and brute-force, so there is a lot of room for improvement
    def propagate_cons(self, cons):
        #print("propagating cons %s with id %d"%(cons.name, id(cons)))
        vars = cons.data.vars
        domains = cons.data.domains

        # TODO: would probably nice to have a flag to know whether the constraint is already propagated
        # I think for this we need and event handler to let us know whenever a variable of our constraint
        # changed its domain

        G, vals = self.build_value_graph(vars, domains)
        M = networkx.bipartite.maximum_matching(G) # returns dict between nodes in matching
        if( len(M)/2 < len(vars) ):
            #print("it is infeasible: max matching of card ", len(M), " M: ", M)
            #print("Its value graph:\nV = ", G.nodes(), "\nE = ", G.edges())
            plot_graph(G)
            return SCIP_RESULT.CUTOFF

        # build auxiliary directed graph: direct var -> val if [var, val] is in matching, otherwise var <- val
        # note that all vars are matched
        D = networkx.DiGraph()
        D.add_nodes_from(G) ## this seems to work
        for var in vars:
            D.add_edge(var.name, M[var.name])
            for val in domains[var]:
                if val != M[var.name]:
                    D.add_edge(val, var.name)

        # find arcs that *do not* need to be removed and *remove* them from G. All remaining edges of G
        # should be use to remove values from the domain of variables
        # get all free vertices
        V = set(G.nodes())
        V_matched = set(M)
        V_free = V.difference(V_matched)
        #print("matched nodes ", V_matched, "\nfree nodes ", V_free)
        # TODO quit() << this produces an assertion

        # no variable should be free!
        for var in vars:
            assert var.name not in V_free

        # perform breadth first search starting from free vertices and mark all visited edges as useful
        for v in V_free:
            visited_edges = bfs_edges(D, v)
            G.remove_edges_from(visited_edges)

        # compute strongly connected components of D and mark edges on the cc as useful
        gscc = networkx.strongly_connected_component_subgraphs(D, copy=False)
        for g in gscc:
            for e in g.edges_iter():
                if G.has_edge(*e):
                    G.remove_edge(*e)

        # cannot remove edges in matching!
        for var in vars:
            e = (var.name, M[var.name])
            if G.has_edge(*e):
                G.remove_edge(*e)

        # check that there is something to remove
        if G.size() == 0:
            return SCIP_RESULT.DIDNOTFIND

        #print("Edges to remove!", G.edges())
        # remove values
        for var in vars:
            for val in domains[var].copy():
                if G.has_edge(var.name, val):
                    domains[var].remove(val) # this asserts if value is not there and we shouldn't delete two times the same value

        # "fix" variable when possible
        for var in vars:
            #print("domain of var ", var.name, "is ", domains[var])
            minval = min(domains[var])
            maxval = max(domains[var])
            if var.getLbLocal() <  minval:
                self.model.chgVarLb(var, minval)
            if var.getUbLocal() > maxval:
                self.model.chgVarUb(var, maxval)
            #print("bounds of ", var, "are (%d,%d)"%(minval,maxval))

        return SCIP_RESULT.REDUCEDDOM

    # propagator callback
    def consprop(self, constraints, nusefulconss, nmarkedconss, proptiming): # I have no idea what to return, documentation?
        result = SCIP_RESULT.DIDNOTFIND
        for cons in constraints:
            prop_result = self.propagate_cons(cons)
            if prop_result == SCIP_RESULT.CUTOFF:
                result = prop_result
                break
            if prop_result == SCIP_RESULT.REDUCEDDOM:
                result = prop_result

        return {"result": result}

    def is_cons_feasible(self, cons, solution = None):
        #print("checking feasibility of constraint %s id: %d"%(cons.name, id(cons)))
        sol_values = set()
        for var in cons.data.vars:
            sol_values.add(round(self.model.getSolVal(solution, var)))
        #print("sol_values = ", sol_values)
        return len(sol_values) == len(cons.data.vars)

    # checks whether solution is feasible, ie, if they are all different
    # since the checkpriority is < 0, we are only called if the integrality
    # constraint handler didn't find infeasibility, so solution is integral
    def conscheck(self, constraints, solution, check_integrality, check_lp_rows, print_reason):
        for cons in constraints:
            if not self.is_cons_feasible(cons, solution):
                return {"result": SCIP_RESULT.INFEASIBLE}
        return {"result": SCIP_RESULT.FEASIBLE}

    # enforces LP solution
    def consenfolp(self, constraints, n_useful_conss, sol_infeasible):
        for cons in constraints:
            if not self.is_cons_feasible(cons):
                # TODO: suggest some value to branch on
                return {"result": SCIP_RESULT.INFEASIBLE}
        return {"result": SCIP_RESULT.FEASIBLE}

    def conslock(self, constraint, nlockspos, nlocksneg):
        for var in constraint.data.vars:
            self.model.addVarLocks(var, nlockspos + nlocksneg , nlockspos + nlocksneg)

    def constrans(self, constraint):
        #print("CONSTRANS BEING CAAAAAAAAAAAAAAAAAAAALLLLLLED")
        return {}

# builds sudoku model; adds variables and all diff constraints
def create_sudoku():
    scip = Model("Sudoku")

    x = {} # values of squares
    for row in range(9):
        for col in range(9):
            # some variables are fix
            if init[row*9 + col] != 0:
                x[row,col] = scip.addVar(vtype = "I", lb = init[row*9 + col], ub = init[row*9 + col], name = "x(%s,%s)" % (row,col))
            else:
                x[row,col] = scip.addVar(vtype = "I", lb = 1, ub = 9, name = "x(%s,%s)" % (row,col))
            var = x[row,col]
            #print("built var ", var.name, " with bounds: (%d,%d)"%(var.getLbLocal(), var.getUbLocal()))

    conshdlr = ALLDIFFconshdlr()

    # hoping to get called when all vars have integer values
    scip.includeConshdlr(conshdlr, "ALLDIFF", "All different constraint", propfreq = 1, enfopriority = -10, chckpriority = -10)

    # row constraints; also we specify the domain of all variables here
    # TODO/QUESTION: in principle domain is of course associated to the var and not the constraint. it should be "var.data"
    # But ideally that information would be handle by SCIP herself... the reason we can't is because domain holes is not implemented, right?
    domains = {}
    for row in range(9):
        vars = []
        for col in range(9):
            var = x[row,col]
            vars.append(var)
            vals = set(range(round(var.getLbLocal()), round(var.getUbLocal()) + 1))
            domains[var] = vals
        # this is kind of ugly, isn't it?
        cons = scip.createCons(conshdlr, "row_%d"%row)
        #print("in test: received a constraint with id ", id(cons)) ### DELETE
        cons.data = SimpleNamespace() # so that data behaves like an instance of a class (ie, cons.data.whatever is allowed)
        cons.data.vars = vars
        cons.data.domains = domains
        scip.addPyCons(cons)

    # col constraints
    for col in range(9):
        vars = []
        for row in range(9):
            var = x[row,col]
            vars.append(var)
        cons = scip.createCons(conshdlr, "col_%d"%col)
        cons.data = SimpleNamespace()
        cons.data.vars = vars
        cons.data.domains = domains
        scip.addPyCons(cons)

    # square constraints
    for idx1 in range(3):
        for idx2 in range(3):
            vars = []
            for row in range(3):
                for col in range(3):
                    var = x[3*idx1 + row, 3*idx2 + col]
                    vars.append(var)
            cons = scip.createCons(conshdlr, "square_%d-%d"%(idx1, idx2))
            cons.data = SimpleNamespace()
            cons.data.vars = vars
            cons.data.domains = domains
            scip.addPyCons(cons)

    #scip.setObjective()

    return scip, x

def test_main():
    scip, x = create_sudoku()

    #scip.setBoolParam("misc/allowdualreds", False)
    scip.setBoolParam("misc/allowdualreds", False)
    scip.setEmphasis(SCIP_PARAMEMPHASIS.CPSOLVER)
    scip.setPresolve(SCIP_PARAMSETTING.OFF)
    scip.optimize()

    if scip.getStatus() != 'optimal':
        print('Sudoku is not feasible!')
    else:
        print('\nSudoku solution:\n')
        for row in range(9):
            out = ''
            for col in range(9):
                out += str(round(scip.getVal(x[row,col]))) + ' '
            print(out)

if __name__ == "__main__":
    test_main()
<commit_msg>fix test_alldiff: use var.ptr() as key since vars are no longer hashable<commit_after>try:
    import networkx
    from networkx.algorithms import bipartite
    from pyscipopt import Model, Conshdlr, SCIP_RESULT, SCIP_PARAMEMPHASIS, SCIP_PARAMSETTING
    from types import SimpleNamespace
    import matplotlib.pyplot as plt
except:
    import pytest
    pytest.skip()

#initial Sudoku values
init = [5, 3, 0, 0, 7, 0, 0, 0, 0,
        6, 0, 0, 1, 9, 5, 0, 0, 0,
        0, 9, 8, 0, 0, 0, 0, 6, 0,
        8, 0, 0, 0, 6, 0, 0, 0, 3,
        4, 0, 0, 8, 0, 3, 0, 0, 1,
        7, 0, 0, 0, 2, 0, 0, 0, 6,
        0, 6, 0, 0, 0, 0, 2, 8, 0,
        0, 0, 0, 4, 1, 9, 0, 0, 5,
        0, 0, 0, 0, 8, 0, 0, 7, 9]

def plot_graph(G):
    X,Y=bipartite.sets(G)
    pos = dict()
    pos.update( (n, (1, i)) for i, n in enumerate(X) ) # put nodes from X at x=1
    pos.update( (n, (2, i)) for i, n in enumerate(Y) ) # put nodes from Y at x=2
    networkx.draw(G, pos=pos, with_labels=False)

    labels = {}
    for node in G.nodes():
        labels[node] = node

    networkx.draw_networkx_labels(G, pos, labels)
    plt.show()

# all different constraint handler
class ALLDIFFconshdlr(Conshdlr):

    # value graph: bipartite graph between variables and the union of their domains
    # an edge connects a variable and a value iff the value is in the variable's domain
    def build_value_graph(self, vars, domains):
        #print(domains)
        vals = set([])
        for var in vars:
            #print("domain of var ", var.name, "is ", domains[var])
            vals.update(domains[var.ptr()]) # vals = vals union domains[var]

        G = networkx.Graph()
        G.add_nodes_from((var.name for var in vars), bipartite = 0) # add vars names as nodes
        G.add_nodes_from(vals, bipartite = 1)                       # add union of values as nodes

        for var in vars:
            for value in domains[var.ptr()]:
                G.add_edge(var.name, value)

        return G, vals

    # propagates single constraint: uses Regin's Algorithm as described in
    # https://www.ps.uni-saarland.de/courses/seminar-ws04/papers/anastasatos.pdf
    # The idea is that every solution of an all different constraint corresponds to a maximal matching in
    # a bipartite graph (see value graph). Furthermore, if an arc of this arc is in no maximal matching, then
    # one can remove it. Removing and arc corresponds to remove a value in the domain of the variable.
    # So what the algorithm does is to determine which arcs can be in a maximal matching. Graph theory help
    # us build fast algorithm so that we don't have to compute all possible maximal matchings ;)
    # That being said, the implementation is pretty naive and brute-force, so there is a lot of room for improvement
    def propagate_cons(self, cons):
        #print("propagating cons %s with id %d"%(cons.name, id(cons)))
        vars = cons.data.vars
        domains = cons.data.domains

        # TODO: would be nice to have a flag to know whether we should propagate the constraint.
        # We would need an event handler to let us know whenever a variable of our constraint changed its domain
        # Currently we can't write event handlers in python.

        G, vals = self.build_value_graph(vars, domains)
        M = networkx.bipartite.maximum_matching(G) # returns dict between nodes in matching
        if( len(M)/2 < len(vars) ):
            #print("it is infeasible: max matching of card ", len(M), " M: ", M)
            #print("Its value graph:\nV = ", G.nodes(), "\nE = ", G.edges())
            plot_graph(G)
            return SCIP_RESULT.CUTOFF

        # build auxiliary directed graph: direct var -> val if [var, val] is in matching, otherwise var <- val
        # note that all vars are matched
        D = networkx.DiGraph()
        D.add_nodes_from(G) ## this seems to work
        for var in vars:
            D.add_edge(var.name, M[var.name])
            for val in domains[var.ptr()]:
                if val != M[var.name]:
                    D.add_edge(val, var.name)

        # find arcs that *do not* need to be removed and *remove* them from G. All remaining edges of G
        # should be use to remove values from the domain of variables
        # get all free vertices
        V = set(G.nodes())
        V_matched = set(M)
        V_free = V.difference(V_matched)
        #print("matched nodes ", V_matched, "\nfree nodes ", V_free)
        # TODO quit() << this produces an assertion

        # no variable should be free!
        for var in vars:
            assert var.name not in V_free

        # perform breadth first search starting from free vertices and mark all visited edges as useful
        for v in V_free:
            visited_edges = bfs_edges(D, v)
            G.remove_edges_from(visited_edges)

        # compute strongly connected components of D and mark edges on the cc as useful
        gscc = networkx.strongly_connected_component_subgraphs(D, copy=False)
        for g in gscc:
            for e in g.edges_iter():
                if G.has_edge(*e):
                    G.remove_edge(*e)

        # cannot remove edges in matching!
        for var in vars:
            e = (var.name, M[var.name])
            if G.has_edge(*e):
                G.remove_edge(*e)

        # check that there is something to remove
        if G.size() == 0:
            return SCIP_RESULT.DIDNOTFIND

        #print("Edges to remove!", G.edges())
        # remove values
        for var in vars:
            for val in domains[var.ptr()].copy():
                if G.has_edge(var.name, val):
                    domains[var.ptr()].remove(val) # this asserts if value is not there and we shouldn't delete two times the same value

        # "fix" variable when possible
        for var in vars:
            #print("domain of var ", var.name, "is ", domains[var])
            minval = min(domains[var.ptr()])
            maxval = max(domains[var.ptr()])
            if var.getLbLocal() <  minval:
                self.model.chgVarLb(var, minval)
            if var.getUbLocal() > maxval:
                self.model.chgVarUb(var, maxval)
            #print("bounds of ", var, "are (%d,%d)"%(minval,maxval))

        return SCIP_RESULT.REDUCEDDOM

    # propagator callback
    def consprop(self, constraints, nusefulconss, nmarkedconss, proptiming): # I have no idea what to return, documentation?
        result = SCIP_RESULT.DIDNOTFIND
        for cons in constraints:
            prop_result = self.propagate_cons(cons)
            if prop_result == SCIP_RESULT.CUTOFF:
                result = prop_result
                break
            if prop_result == SCIP_RESULT.REDUCEDDOM:
                result = prop_result

        return {"result": result}

    def is_cons_feasible(self, cons, solution = None):
        #print("checking feasibility of constraint %s id: %d"%(cons.name, id(cons)))
        sol_values = set()
        for var in cons.data.vars:
            sol_values.add(round(self.model.getSolVal(solution, var)))
        #print("sol_values = ", sol_values)
        return len(sol_values) == len(cons.data.vars)

    # checks whether solution is feasible, ie, if they are all different
    # since the checkpriority is < 0, we are only called if the integrality
    # constraint handler didn't find infeasibility, so solution is integral
    def conscheck(self, constraints, solution, check_integrality, check_lp_rows, print_reason):
        for cons in constraints:
            if not self.is_cons_feasible(cons, solution):
                return {"result": SCIP_RESULT.INFEASIBLE}
        return {"result": SCIP_RESULT.FEASIBLE}

    # enforces LP solution
    def consenfolp(self, constraints, n_useful_conss, sol_infeasible):
        for cons in constraints:
            if not self.is_cons_feasible(cons):
                # TODO: suggest some value to branch on
                return {"result": SCIP_RESULT.INFEASIBLE}
        return {"result": SCIP_RESULT.FEASIBLE}

    def conslock(self, constraint, nlockspos, nlocksneg):
        for var in constraint.data.vars:
            self.model.addVarLocks(var, nlockspos + nlocksneg , nlockspos + nlocksneg)

    def constrans(self, constraint):
        #print("CONSTRANS BEING CAAAAAAAAAAAAAAAAAAAALLLLLLED")
        return {}

# builds sudoku model; adds variables and all diff constraints
def create_sudoku():
    scip = Model("Sudoku")

    x = {} # values of squares
    for row in range(9):
        for col in range(9):
            # some variables are fix
            if init[row*9 + col] != 0:
                x[row,col] = scip.addVar(vtype = "I", lb = init[row*9 + col], ub = init[row*9 + col], name = "x(%s,%s)" % (row,col))
            else:
                x[row,col] = scip.addVar(vtype = "I", lb = 1, ub = 9, name = "x(%s,%s)" % (row,col))
            var = x[row,col]
            #print("built var ", var.name, " with bounds: (%d,%d)"%(var.getLbLocal(), var.getUbLocal()))

    conshdlr = ALLDIFFconshdlr()

    # hoping to get called when all vars have integer values
    scip.includeConshdlr(conshdlr, "ALLDIFF", "All different constraint", propfreq = 1, enfopriority = -10, chckpriority = -10)

    # row constraints; also we specify the domain of all variables here
    # TODO/QUESTION: in principle domain is of course associated to the var and not the constraint. it should be "var.data"
    # But ideally that information would be handle by SCIP itself... the reason we can't is because domain holes is not implemented, right?
    domains = {}
    for row in range(9):
        vars = []
        for col in range(9):
            var = x[row,col]
            vars.append(var)
            vals = set(range(round(var.getLbLocal()), round(var.getUbLocal()) + 1))
            domains[var.ptr()] = vals
        # this is kind of ugly, isn't it?
        cons = scip.createCons(conshdlr, "row_%d"%row)
        #print("in test: received a constraint with id ", id(cons)) ### DELETE
        cons.data = SimpleNamespace() # so that data behaves like an instance of a class (ie, cons.data.whatever is allowed)
        cons.data.vars = vars
        cons.data.domains = domains
        scip.addPyCons(cons)

    # col constraints
    for col in range(9):
        vars = []
        for row in range(9):
            var = x[row,col]
            vars.append(var)
        cons = scip.createCons(conshdlr, "col_%d"%col)
        cons.data = SimpleNamespace()
        cons.data.vars = vars
        cons.data.domains = domains
        scip.addPyCons(cons)

    # square constraints
    for idx1 in range(3):
        for idx2 in range(3):
            vars = []
            for row in range(3):
                for col in range(3):
                    var = x[3*idx1 + row, 3*idx2 + col]
                    vars.append(var)
            cons = scip.createCons(conshdlr, "square_%d-%d"%(idx1, idx2))
            cons.data = SimpleNamespace()
            cons.data.vars = vars
            cons.data.domains = domains
            scip.addPyCons(cons)

    #scip.setObjective()

    return scip, x

def test_main():
    scip, x = create_sudoku()

    #scip.setBoolParam("misc/allowdualreds", False)
    scip.setBoolParam("misc/allowdualreds", False)
    scip.setEmphasis(SCIP_PARAMEMPHASIS.CPSOLVER)
    scip.setPresolve(SCIP_PARAMSETTING.OFF)
    scip.optimize()

    if scip.getStatus() != 'optimal':
        print('Sudoku is not feasible!')
    else:
        print('\nSudoku solution:\n')
        for row in range(9):
            out = ''
            for col in range(9):
                out += str(round(scip.getVal(x[row,col]))) + ' '
            print(out)

if __name__ == "__main__":
    test_main()
<|endoftext|>


<commit_before>import os
import subprocess
import unittest

from bold_retriever import bold_retriever as br

class CmdlineTest(unittest.TestCase):

    def setUp(self):
        self.parser = br.create_parser()

    # User does not enter any argument
    def test_cmdline_emtpy_args(self):
        with self.assertRaises(SystemExit):
            self.parser.parse_args([])

    # User enters only one argument
    def test_cmdline_file(self):
        with self.assertRaises(SystemExit):
            self.parser.parse_args(['-f', 'ionx23b.fas'])

    # User enters both arguments
    def test_cmdline(self):
        args = self.parser.parse_args(['-f', 'ionx23b.fas', '-db', 'COX1'])
        expected = 'COX1'
        self.assertEqual(expected, args.db)

    def test_cmdline_main(self):
        cwd = os.path.abspath(os.path.dirname(__file__))
        bold_retriever = os.path.join(cwd, '..', 'bold_retriever',
                                      'bold_retriever.py')
        print(bold_retriever)
        cmd = "python " + bold_retriever + " -f " + os.path.join(cwd,
                                                                 'otu99.fas')
        cmd += " -db COX1_SPECIES"
        p = subprocess.check_call(cmd, shell=True)
        self.assertEqual(0, p)

    def test_get_args(self):
        args = self.parser.parse_args(['-f', 'ionx13.fas', '-db', 'COX1'])
        f, db = br.get_args(args)
        self.assertEqual(f, 'ionx13.fas')
        self.assertEqual(db, 'COX1')

    def test_get_started(self):
        test_folder = os.path.abspath(os.path.dirname(__file__))
        input_file = os.path.join(test_folder, 'ionx13.fas')
        result_file = os.path.join(test_folder, 'ionx13.fas_output.csv')

        args = self.parser.parse_args(['-f', input_file, '-db', 'COX1'])
        br.get_started(args)

        self.assertTrue(os.path.isfile(result_file))

    def test_otu99(self):
        """It should pass despite returning empty list for classification."""
        test_folder = os.path.abspath(os.path.dirname(__file__))
        input_file = os.path.join(test_folder, 'otu99.fas')
        args = self.parser.parse_args(['-f', input_file, '-db',
                                       'COX1_SPECIES'])
        br.get_started(args)
        result_file = os.path.join(test_folder, 'otu99.fas_output.csv')
        self.assertTrue(os.path.isfile(result_file))

    def test_insecta(self):
        """It should pass despite returning empty list for classification."""
        test_folder = os.path.abspath(os.path.dirname(__file__))
        input_file = os.path.join(test_folder, 'otu_insecta.fas')
        args = self.parser.parse_args(['-f', input_file, '-db',
                                       'COX1_SPECIES'])
        br.get_started(args)
        result_file = os.path.join(test_folder, 'otu_insecta.fas_output.csv')
        self.assertTrue(os.path.isfile(result_file))
<commit_msg>fix test cmdline<commit_after>import glob
import os
import subprocess
import unittest

from bold_retriever import bold_retriever as br

class CmdlineTest(unittest.TestCase):

    def setUp(self):
        self.parser = br.create_parser()

    # User does not enter any argument
    def test_cmdline_emtpy_args(self):
        with self.assertRaises(SystemExit):
            self.parser.parse_args([])

    # User enters only one argument
    def test_cmdline_file(self):
        with self.assertRaises(SystemExit):
            self.parser.parse_args(['-f', 'ionx23b.fas'])

    # User enters both arguments
    def test_cmdline(self):
        args = self.parser.parse_args(['-f', 'ionx23b.fas', '-db', 'COX1'])
        expected = 'COX1'
        self.assertEqual(expected, args.db)

    def test_cmdline_main(self):
        cwd = os.path.abspath(os.path.dirname(__file__))
        bold_retriever = os.path.join(cwd, '..', 'bold_retriever',
                                      'bold_retriever.py')
        print(bold_retriever)
        cmd = "python " + bold_retriever + " -f " + os.path.join(cwd,
                                                                 'otu99.fas')
        cmd += " -db COX1_SPECIES"
        p = subprocess.check_call(cmd, shell=True)
        self.assertEqual(0, p)

    def test_get_args(self):
        args = self.parser.parse_args(['-f', 'ionx13.fas', '-db', 'COX1'])
        f, db = br.get_args(args)
        self.assertEqual(f, 'ionx13.fas')
        self.assertEqual(db, 'COX1')

    def test_get_started(self):
        test_folder = os.path.abspath(os.path.dirname(__file__))
        input_file = os.path.join(test_folder, 'ionx13.fas')
        result_file = os.path.join(test_folder, 'ionx13.fas_output.csv')

        args = self.parser.parse_args(['-f', input_file, '-db', 'COX1'])
        br.get_started(args)

        self.assertTrue(os.path.isfile(result_file))

    def test_otu99(self):
        """It should pass despite returning empty list for classification."""
        test_folder = os.path.abspath(os.path.dirname(__file__))
        input_file = os.path.join(test_folder, 'otu99.fas')
        args = self.parser.parse_args(['-f', input_file, '-db',
                                       'COX1_SPECIES'])
        br.get_started(args)
        result_file = os.path.join(test_folder, 'otu99.fas_output.csv')
        self.assertTrue(os.path.isfile(result_file))

    def test_insecta(self):
        """It should pass despite returning empty list for classification, insecta."""
        test_folder = os.path.abspath(os.path.dirname(__file__))
        input_file = os.path.join(test_folder, 'otu_insecta.fas')
        args = self.parser.parse_args(['-f', input_file, '-db',
                                       'COX1_SPECIES'])
        br.get_started(args)
        result_file = os.path.join(test_folder, 'otu_insecta.fas_output.csv')
        self.assertTrue(os.path.isfile(result_file))

    def tearDown(self):
        test_folder = os.path.abspath(os.path.dirname(__file__))
        for i in glob.glob(test_folder + "/*csv"):
            os.remove(i)
<|endoftext|>