<commit_before>#
# linter.py
# Part of SublimeLinter3, a code checking framework for Sublime Text 3
#
# Written by Ryan Hileman and Aparajita Fishman
#
# Project: https://github.com/SublimeLinter/SublimeLinter3
# License: MIT
#

import re
import shlex
import sublime

from . import highlight as hilite
from . import persist
from . import util

SYNTAX_RE = re.compile(r'/([^/]+)\.tmLanguage$')
WARNING_RE = re.compile(r'^w(?:arn(?:ing)?)?$', re.IGNORECASE)

class Registrar(type):
    '''This metaclass registers the linter when the class is declared.'''
    def __init__(cls, name, bases, attrs):
        if bases:
            persist.register_linter(cls, name, attrs)

class Linter(metaclass=Registrar):
    '''
    The base class for linters. Subclasses must at a minimum define
    the attributes language, cmd, and regex.
    '''

    #
    # Public attributes
    #

    # The name of the linter's language for display purposes.
    # By convention this is all lowercase.
    language = ''

    # A string, tuple or callable that returns a string or tuple, containing the
    # command line arguments used to lint.
    cmd = ''

    # If the name of the executable cannot be determined by the first element of cmd
    # (for example when cmd is a method that dynamically generates the command line arguments),
    # this can be set to the name of the executable used to do linting.
    #
    # Once the executable's name is determined, its existence is checked in the user's path.
    # If it is not available, the linter is disabled.
    executable = None

    # If the executable is available, this is set to the full path of the executable.
    # If the executable is not available, it is set an empty string.
    # Subclasses should consider this read only.
    executable_path = None

    # A regex pattern used to extract information from the linter's executable output.
    regex = ''

    # Set to True if the linter outputs multiple errors or multiline errors. When True,
    # regex will be created with the re.MULTILINE flag.
    multiline = False

    # If you want to set flags on the regex other than re.MULTILINE, set this.
    re_flags = 0

    # If the linter executable cannot receive from stdin and requires a temp file,
    # set this attribute to the suffix of the temp file (including leading '.').
    tempfile_suffix = None

    # Tab width
    tab_width = 1

    # If you want to limit the linter to specific portions of the source
    # based on a scope selector, set this attribute to the selector. For example,
    # in an html file with embedded php, you would set the selector for a php
    # linter to 'source.php'.
    selector = None

    # If you want to provide default settings for the linter, set this attribute.
    defaults = None

    #
    # Internal class storage, do not set
    #
    errors = None
    highlight = None
    lint_settings = None

    def __init__(self, view, syntax, filename=None):
        self.view = view
        self.syntax = syntax
        self.filename = filename

        if self.regex:
            if self.multiline:
                self.re_flags |= re.MULTILINE

            try:
                self.regex = re.compile(self.regex, self.re_flags)
            except:
                persist.debug('error compiling regex for {}'.format(self.language))

        self.highlight = hilite.Highlight()

    @classmethod
    def get_settings(cls):
        '''Return the default settings for this linter, merged with the user settings.'''
        linters = persist.settings.get('linters', {})
        settings = cls.defaults or {}
        settings.update(linters.get(cls.__name__, {}))
        return settings

    @property
    def settings(self):
        return self.get_settings()

    @classmethod
    def assign(cls, view, reassign=False):
        '''
        Assign a view to an instance of a linter.
        Find a linter for a specified view if possible, then add it to our view <--> lint class map and return it.
        Each view has its own linter so that linters can store persistent data about a view.
        '''
        vid = view.id()
        persist.views[vid] = view

        settings = view.settings()
        syntax = settings.get('syntax')

        if not syntax:
            cls.remove(vid)
            return

        match = SYNTAX_RE.search(syntax)

        if match:
            syntax = match.group(1)
        else:
            syntax = ''

        if syntax:
            if vid in persist.linters and persist.linters[vid] and not reassign:
                # If a linter in the set of linters for the given view
                # already handles the view's syntax, we have nothing more to do.
                for linter in tuple(persist.linters[vid]):
                    if linter.syntax == syntax:
                        return

            linters = set()

            for name, linter_class in persist.languages.items():
                if linter_class.can_lint(syntax):
                    linter = linter_class(view, syntax, view.file_name())
                    linters.add(linter)

            if linters:
                persist.linters[vid] = linters

            return linters

        cls.remove(vid)

    @classmethod
    def remove(cls, vid):
        '''Remove a the mapping between a view and its set of linters.'''
        if vid in persist.linters:
            for linters in persist.linters[vid]:
                linters.clear()

            del persist.linters[vid]

    @classmethod
    def reload(cls, mod=None):
        '''Reload all linters, optionally filtering by module.'''

        # Merge linter default settings with user settings
        linter_settings = persist.settings.get('linters', {})

        for name, linter in persist.languages.items():
            settings = linter_settings.get(name, {})
            defaults = (linter.defaults or {}).copy()
            defaults.update(settings)
            linter.lint_settings = defaults

        for vid, linters in persist.linters.items():
            for linter in linters:
                if mod and linter.__module__ != mod:
                    continue

                linter.clear()
                persist.linters[vid].remove(linter)
                linter = persist.languages[linter.name](linter.view, linter.syntax, linter.filename)
                persist.linters[vid].add(linter)

        cls.redraw_all()

    @classmethod
    def apply_to_all(cls, action):
        for w in sublime.windows():
            for view in w.views():
                highlights = persist.highlights.get(view.id())

                if highlights:
                    getattr(highlights, action)(view)

    @classmethod
    def clear_all(cls):
        cls.apply_to_all('reset')
        persist.errors.clear()

    @classmethod
    def redraw_all(cls):
        cls.apply_to_all('redraw')

    @classmethod
    def text(cls, view):
        '''Returns the entire text of a view.'''
        return view.substr(sublime.Region(0, view.size()))

    @classmethod
    def get_view(cls, vid):
        '''Returns the view object with the given id.'''
        return persist.views.get(vid)

    @classmethod
    def get_linters(cls, vid):
        '''Returns a tuple of linters for the view with the given id.'''
        if vid in persist.linters:
            return tuple(persist.linters[vid])

        return ()

    @classmethod
    def get_selectors(cls, vid):
        '''Returns a list of scope selectors for all linters for the view with the given id.'''
        return [
            (linter.selector, linter)
            for linter in cls.get_linters(vid)
            if linter.selector
        ]

    @classmethod
    def lint_view(cls, vid, filename, code, sections, hit_time, callback):
        if not code or vid not in persist.linters:
            return

        linters = list(persist.linters.get(vid))

        if not linters:
            return

        filename = filename or 'untitled'

        for linter in linters:
            if linter.settings.get('disable'):
                continue

            if not linter.selector:
                linter.reset(code, filename=filename)
                linter.lint()

        selectors = Linter.get_selectors(vid)

        for sel, linter in selectors:
            linters.append(linter)

            if sel in sections:
                linter.reset(code, filename=filename)
                errors = {}

                for line_offset, left, right in sections[sel]:
                    linter.hilite.move_to(line_offset, left)
                    linter.code = code[left:right]
                    linter.errors = {}
                    linter.lint()

                    for line, error in linter.errors.items():
                        errors[line + line_offset] = error

                linter.errors = errors

        # Merge our result back to the main thread
        callback(cls.get_view(vid), linters, hit_time)

    def reset(self, code, filename=None, highlight=None):
        self.errors = {}
        self.code = code
        self.filename = filename or self.filename
        self.highlight = highlight or hilite.Highlight(self.code)

    def get_cmd(self):
        if callable(self.cmd):
            cmd = self.cmd()
        else:
            cmd = self.cmd

        if not cmd:
            return

        if isinstance(cmd, str):
            cmd = shlex.split(cmd)

        return tuple(cmd)

    def lint(self):
        if not (self.language and self.cmd and self.regex):
            raise NotImplementedError

        cmd = self.get_cmd()

        if not cmd:
            return

        output = self.run(cmd, self.code)

        if not output:
            return

        persist.debug('{} output:\n{}'.format(self.__class__.__name__, output.strip()))

        for match, row, col, error_type, message, near in self.find_errors(output):
            if match and row is not None:
                if error_type and WARNING_RE.match(error_type) is not None:
                    error_type = hilite.WARNING
                else:
                    error_type = hilite.ERROR

                if col is not None:
                    # Adjust column numbers to match the linter's tabs if necessary
                    if self.tab_width > 1:
                        start, end = self.highlight.full_line(row)
                        code_line = self.code[start:end]
                        diff = 0

                        for i in range(len(code_line)):
                            if code_line[i] == '\t':
                                diff += (self.tab_width - 1)

                            if col - diff <= i:
                                col = i
                                break

                    self.highlight.range(row, col, error_type=error_type)
                elif near:
                    col = self.highlight.near(row, near, error_type)
                else:
                    self.highlight.range(row, 0, length=0, error_type=error_type)

                self.error(row, col, message, error_type)

    def draw(self):
        self.highlight.draw(self.view)

    def clear(self):
        self.highlight.clear(self.view)

    # Helper methods

    @classmethod
    def can_lint(cls, language):
        '''
        Determines if a linter can lint a given language. Subclasses may override this
        if the built in mechanism is not sufficient, but should call super().can_lint(cls, language)
        first and continue checking only if that returns True.
        '''
        can = False
        language = language.lower()

        if cls.language:
            if language == cls.language:
                can = True
            elif isinstance(cls.language, (tuple, list)) and language in cls.language:
                can = True

        if can and cls.executable_path is None:
            executable = ''

            if not callable(cls.cmd):
                if isinstance(cls.cmd, (tuple, list)):
                    executable = (cls.cmd or [''])[0]
                else:
                    executable = cls.cmd

            if not executable and cls.executable:
                executable = cls.executable

            if executable:
                cls.executable_path = util.which(executable) or ''
            else:
                cls.executable_path = ''

            can = cls.executable_path != ''
            persist.printf('{} {}'.format(
                cls.__name__,
                'enabled ({})'.format(cls.executable_path) if can
                else 'disabled, cannot locate \'{}\''.format(executable)
            ))

        return can

    def error(self, line, col, error, error_type):
        self.highlight.line(line, error_type)

        # Capitalize the first word
        error = error[0].upper() + error[1:]

        # Strip trailing space and period
        error = ((col or 0), str(error).rstrip(' .'))

        if line in self.errors:
            self.errors[line].append(error)
        else:
            self.errors[line] = [error]

    def find_errors(self, output):
        if self.multiline:
            errors = self.regex.finditer(output)

            if errors:
                for error in errors:
                    yield self.split_match(error)
            else:
                yield self.split_match(None)
        else:
            for line in output.splitlines():
                yield self.match_error(self.regex, line.strip())

    def split_match(self, match):
        if match:
            items = {'line': None, 'col': None, 'type': None, 'error': '', 'near': None}
            items.update(match.groupdict())
            row, col, error_type, error, near = [items[k] for k in ('line', 'col', 'type', 'error', 'near')]

            if row is not None:
                row = int(row) - 1

            if col is not None:
                col = int(col) - 1

            return match, row, col, error_type, error, near
        else:
            return match, None, None, None, None, '', None

    def match_error(self, r, line):
        return self.split_match(r.match(line))

    # Subclasses may need to override this in complex cases
    def run(self, cmd, code):
        if self.tempfile_suffix:
            return self.tmpfile(cmd, code, suffix=self.tempfile_suffix)
        else:
            return self.communicate(cmd, code)

    # popen wrappers
    def communicate(self, cmd, code):
        return util.communicate(cmd, code)

    def tmpfile(self, cmd, code, suffix=''):
        return util.tmpfile(cmd, code, suffix or self.tempfile_suffix)

    def tmpdir(self, cmd, files, code):
        return util.tmpdir(cmd, files, self.filename, code)

    def popen(self, cmd, env=None):
        return util.popen(cmd, env)
<commit_msg>Added a subclass hook for determining if a linter can lint a language<commit_after>#
# linter.py
# Part of SublimeLinter3, a code checking framework for Sublime Text 3
#
# Written by Ryan Hileman and Aparajita Fishman
#
# Project: https://github.com/SublimeLinter/SublimeLinter3
# License: MIT
#

import re
import shlex
import sublime

from . import highlight as hilite
from . import persist
from . import util

SYNTAX_RE = re.compile(r'/([^/]+)\.tmLanguage$')
WARNING_RE = re.compile(r'^w(?:arn(?:ing)?)?$', re.IGNORECASE)

class Registrar(type):
    '''This metaclass registers the linter when the class is declared.'''
    def __init__(cls, name, bases, attrs):
        if bases:
            persist.register_linter(cls, name, attrs)

class Linter(metaclass=Registrar):
    '''
    The base class for linters. Subclasses must at a minimum define
    the attributes language, cmd, and regex.
    '''

    #
    # Public attributes
    #

    # The name of the linter's language for display purposes.
    # By convention this is all lowercase.
    language = ''

    # A string, tuple or callable that returns a string or tuple, containing the
    # command line arguments used to lint.
    cmd = ''

    # If the name of the executable cannot be determined by the first element of cmd
    # (for example when cmd is a method that dynamically generates the command line arguments),
    # this can be set to the name of the executable used to do linting.
    #
    # Once the executable's name is determined, its existence is checked in the user's path.
    # If it is not available, the linter is disabled.
    executable = None

    # If the executable is available, this is set to the full path of the executable.
    # If the executable is not available, it is set an empty string.
    # Subclasses should consider this read only.
    executable_path = None

    # A regex pattern used to extract information from the linter's executable output.
    regex = ''

    # Set to True if the linter outputs multiple errors or multiline errors. When True,
    # regex will be created with the re.MULTILINE flag.
    multiline = False

    # If you want to set flags on the regex other than re.MULTILINE, set this.
    re_flags = 0

    # If the linter executable cannot receive from stdin and requires a temp file,
    # set this attribute to the suffix of the temp file (including leading '.').
    tempfile_suffix = None

    # Tab width
    tab_width = 1

    # If you want to limit the linter to specific portions of the source
    # based on a scope selector, set this attribute to the selector. For example,
    # in an html file with embedded php, you would set the selector for a php
    # linter to 'source.php'.
    selector = None

    # If you want to provide default settings for the linter, set this attribute.
    defaults = None

    #
    # Internal class storage, do not set
    #
    errors = None
    highlight = None
    lint_settings = None

    def __init__(self, view, syntax, filename=None):
        self.view = view
        self.syntax = syntax
        self.filename = filename

        if self.regex:
            if self.multiline:
                self.re_flags |= re.MULTILINE

            try:
                self.regex = re.compile(self.regex, self.re_flags)
            except:
                persist.debug('error compiling regex for {}'.format(self.language))

        self.highlight = hilite.Highlight()

    @classmethod
    def get_settings(cls):
        '''Return the default settings for this linter, merged with the user settings.'''
        linters = persist.settings.get('linters', {})
        settings = cls.defaults or {}
        settings.update(linters.get(cls.__name__, {}))
        return settings

    @property
    def settings(self):
        return self.get_settings()

    @classmethod
    def assign(cls, view, reassign=False):
        '''
        Assign a view to an instance of a linter.
        Find a linter for a specified view if possible, then add it to our view <--> lint class map and return it.
        Each view has its own linter so that linters can store persistent data about a view.
        '''
        vid = view.id()
        persist.views[vid] = view

        settings = view.settings()
        syntax = settings.get('syntax')

        if not syntax:
            cls.remove(vid)
            return

        match = SYNTAX_RE.search(syntax)

        if match:
            syntax = match.group(1)
        else:
            syntax = ''

        if syntax:
            if vid in persist.linters and persist.linters[vid] and not reassign:
                # If a linter in the set of linters for the given view
                # already handles the view's syntax, we have nothing more to do.
                for linter in tuple(persist.linters[vid]):
                    if linter.syntax == syntax:
                        return

            linters = set()

            for name, linter_class in persist.languages.items():
                if linter_class.can_lint(syntax):
                    linter = linter_class(view, syntax, view.file_name())
                    linters.add(linter)

            if linters:
                persist.linters[vid] = linters

            return linters

        cls.remove(vid)

    @classmethod
    def remove(cls, vid):
        '''Remove a the mapping between a view and its set of linters.'''
        if vid in persist.linters:
            for linters in persist.linters[vid]:
                linters.clear()

            del persist.linters[vid]

    @classmethod
    def reload(cls, mod=None):
        '''Reload all linters, optionally filtering by module.'''

        # Merge linter default settings with user settings
        linter_settings = persist.settings.get('linters', {})

        for name, linter in persist.languages.items():
            settings = linter_settings.get(name, {})
            defaults = (linter.defaults or {}).copy()
            defaults.update(settings)
            linter.lint_settings = defaults

        for vid, linters in persist.linters.items():
            for linter in linters:
                if mod and linter.__module__ != mod:
                    continue

                linter.clear()
                persist.linters[vid].remove(linter)
                linter = persist.languages[linter.name](linter.view, linter.syntax, linter.filename)
                persist.linters[vid].add(linter)

        cls.redraw_all()

    @classmethod
    def apply_to_all(cls, action):
        for w in sublime.windows():
            for view in w.views():
                highlights = persist.highlights.get(view.id())

                if highlights:
                    getattr(highlights, action)(view)

    @classmethod
    def clear_all(cls):
        cls.apply_to_all('reset')
        persist.errors.clear()

    @classmethod
    def redraw_all(cls):
        cls.apply_to_all('redraw')

    @classmethod
    def text(cls, view):
        '''Returns the entire text of a view.'''
        return view.substr(sublime.Region(0, view.size()))

    @classmethod
    def get_view(cls, vid):
        '''Returns the view object with the given id.'''
        return persist.views.get(vid)

    @classmethod
    def get_linters(cls, vid):
        '''Returns a tuple of linters for the view with the given id.'''
        if vid in persist.linters:
            return tuple(persist.linters[vid])

        return ()

    @classmethod
    def get_selectors(cls, vid):
        '''Returns a list of scope selectors for all linters for the view with the given id.'''
        return [
            (linter.selector, linter)
            for linter in cls.get_linters(vid)
            if linter.selector
        ]

    @classmethod
    def lint_view(cls, vid, filename, code, sections, hit_time, callback):
        if not code or vid not in persist.linters:
            return

        linters = list(persist.linters.get(vid))

        if not linters:
            return

        filename = filename or 'untitled'

        for linter in linters:
            if linter.settings.get('disable'):
                continue

            if not linter.selector:
                linter.reset(code, filename=filename)
                linter.lint()

        selectors = Linter.get_selectors(vid)

        for sel, linter in selectors:
            linters.append(linter)

            if sel in sections:
                linter.reset(code, filename=filename)
                errors = {}

                for line_offset, left, right in sections[sel]:
                    linter.hilite.move_to(line_offset, left)
                    linter.code = code[left:right]
                    linter.errors = {}
                    linter.lint()

                    for line, error in linter.errors.items():
                        errors[line + line_offset] = error

                linter.errors = errors

        # Merge our result back to the main thread
        callback(cls.get_view(vid), linters, hit_time)

    def reset(self, code, filename=None, highlight=None):
        self.errors = {}
        self.code = code
        self.filename = filename or self.filename
        self.highlight = highlight or hilite.Highlight(self.code)

    def get_cmd(self):
        if callable(self.cmd):
            cmd = self.cmd()
        else:
            cmd = self.cmd

        if not cmd:
            return

        if isinstance(cmd, str):
            cmd = shlex.split(cmd)

        return tuple(cmd)

    def lint(self):
        if not (self.language and self.cmd and self.regex):
            raise NotImplementedError

        cmd = self.get_cmd()

        if not cmd:
            return

        output = self.run(cmd, self.code)

        if not output:
            return

        persist.debug('{} output:\n{}'.format(self.__class__.__name__, output.strip()))

        for match, row, col, error_type, message, near in self.find_errors(output):
            if match and row is not None:
                if error_type and WARNING_RE.match(error_type) is not None:
                    error_type = hilite.WARNING
                else:
                    error_type = hilite.ERROR

                if col is not None:
                    # Adjust column numbers to match the linter's tabs if necessary
                    if self.tab_width > 1:
                        start, end = self.highlight.full_line(row)
                        code_line = self.code[start:end]
                        diff = 0

                        for i in range(len(code_line)):
                            if code_line[i] == '\t':
                                diff += (self.tab_width - 1)

                            if col - diff <= i:
                                col = i
                                break

                    self.highlight.range(row, col, error_type=error_type)
                elif near:
                    col = self.highlight.near(row, near, error_type)
                else:
                    self.highlight.range(row, 0, length=0, error_type=error_type)

                self.error(row, col, message, error_type)

    def draw(self):
        self.highlight.draw(self.view)

    def clear(self):
        self.highlight.clear(self.view)

    # Helper methods

    @classmethod
    def can_lint(cls, language):
        can = False
        language = language.lower()

        if cls.language:
            if isinstance(cls.language, (tuple, list)) and language in cls.language:
                can = True
            elif language == cls.language:
                can = True

        if can and cls.executable_path is None:
            executable = ''

            if not callable(cls.cmd):
                if isinstance(cls.cmd, (tuple, list)):
                    executable = (cls.cmd or [''])[0]
                else:
                    executable = cls.cmd

            if not executable and cls.executable:
                executable = cls.executable

            if executable:
                cls.executable_path = util.which(executable) or ''
            else:
                cls.executable_path = ''

            can = cls.executable_path != ''

        if can:
            can = cls.can_lint_language(language)

        persist.printf('{} {}'.format(
            cls.__name__,
            'enabled ({})'.format(cls.executable_path) if can
            else 'disabled, cannot locate \'{}\''.format(executable)
        ))

        return can

    @classmethod
    def can_lint_language(cls, language):
        '''
        Subclass hook to determine if a linter can lint a given language. Subclasses may override this
        if the built in mechanism is not sufficient. Note that this will only be called if the built
        in checks succeed first. When this is called, cls.executable_path has been set.
        '''
        return True

    def error(self, line, col, error, error_type):
        self.highlight.line(line, error_type)

        # Capitalize the first word
        error = error[0].upper() + error[1:]

        # Strip trailing space and period
        error = ((col or 0), str(error).rstrip(' .'))

        if line in self.errors:
            self.errors[line].append(error)
        else:
            self.errors[line] = [error]

    def find_errors(self, output):
        if self.multiline:
            errors = self.regex.finditer(output)

            if errors:
                for error in errors:
                    yield self.split_match(error)
            else:
                yield self.split_match(None)
        else:
            for line in output.splitlines():
                yield self.match_error(self.regex, line.strip())

    def split_match(self, match):
        if match:
            items = {'line': None, 'col': None, 'type': None, 'error': '', 'near': None}
            items.update(match.groupdict())
            row, col, error_type, error, near = [items[k] for k in ('line', 'col', 'type', 'error', 'near')]

            if row is not None:
                row = int(row) - 1

            if col is not None:
                col = int(col) - 1

            return match, row, col, error_type, error, near
        else:
            return match, None, None, None, None, '', None

    def match_error(self, r, line):
        return self.split_match(r.match(line))

    # Subclasses may need to override this in complex cases
    def run(self, cmd, code):
        if self.tempfile_suffix:
            return self.tmpfile(cmd, code, suffix=self.tempfile_suffix)
        else:
            return self.communicate(cmd, code)

    # popen wrappers
    def communicate(self, cmd, code):
        return util.communicate(cmd, code)

    def tmpfile(self, cmd, code, suffix=''):
        return util.tmpfile(cmd, code, suffix or self.tempfile_suffix)

    def tmpdir(self, cmd, files, code):
        return util.tmpdir(cmd, files, self.filename, code)

    def popen(self, cmd, env=None):
        return util.popen(cmd, env)
<|endoftext|>


<commit_before>import unittest
import functools

import numpy

import cupy
from cupy import testing
from cupyx import fallback_mode

@testing.gpu
class TestFallbackMode(unittest.TestCase):

    def numpy_fallback_equal(name='xp'):
        """
        Decorator that checks fallback_mode results are equal to NumPy ones.

        Args:
            name(str): Argument name whose value is either
            ``numpy`` or ``cupy`` module.
        """
        def decorator(impl):
            @functools.wraps(impl)
            def test_func(self, *args, **kwargs):

                kwargs[name] = fallback_mode.numpy
                fallback_result = impl(self, *args, **kwargs)

                kwargs[name] = numpy
                numpy_result = impl(self, *args, **kwargs)

                if isinstance(numpy_result, numpy.ndarray):
                    # if numpy returns ndarray, cupy must return ndarray
                    assert isinstance(fallback_result, cupy.ndarray)
                    testing.assert_array_equal(numpy_result, fallback_result)

                elif isinstance(numpy_result, numpy.ScalarType):
                    # if numpy returns scalar
                    # cupy must return scalar or 0-dim array
                    if isinstance(fallback_result, numpy.ScalarType):
                        assert numpy_result == fallback_result

                    else:
                        # cupy 0-dim array
                        assert numpy_result == int(fallback_result)
                else:
                    raise NotImplementedError

            return test_func
        return decorator

    @numpy_fallback_equal()
    def test_argmin(self, xp):

        a = xp.array([
            [13, 5, 45, 23, 9],
            [-5, 41, 0, 22, 4],
            [2, 6, 43, 11, -1]
        ])

        return xp.argmin(a, axis=1)

    @numpy_fallback_equal()
    def test_argmin_zero_dim_array_vs_scalar(self, xp):

        a = xp.array([
            [13, 5, 45, 23, 9],
            [-5, 41, 0, 22, 4],
            [2, 6, 43, 11, -1]
        ])

        return xp.argmin(a)

    # cupy.argmin raises error if list passed, numpy does not
    def test_argmin_list(self):

        a = [
            [13, 5, 45, 23, 9],
            [-5, 41, 0, 22, 4],
            [2, 6, 43, 11, -1]
        ]

        with self.assertRaises(AttributeError):
            fallback_mode.numpy.argmin(a)

        assert numpy.argmin([1, 0, 3]) == 1

    # Non-existing function
    @numpy_fallback_equal()
    def test_array_equal(self, xp):

        a = xp.array([1, 2])
        b = xp.array([1, 2])

        return xp.array_equal(a, b)

    # Both cupy and numpy return 0-d array
    @numpy_fallback_equal()
    def test_convolve_zero_dim_array(self, xp):

        a = xp.array([1, 2, 3])
        b = xp.array([0, 1, 0.5])

        return xp.convolve(a, b, 'valid')

    def test_vectorize(self):

        def function(a, b):
            if a > b:
                return a - b
            return a + b

        actual = numpy.vectorize(function)([1, 2, 3, 4], 2)
        expected = fallback_mode.numpy.vectorize(function)([1, 2, 3, 4], 2)

        assert isinstance(actual, numpy.ndarray)

        # ([1,2,3,4], 2) are arguments to
        # numpy.vectorize(function), not numpy.vectorize
        # So, it returns numpy.ndarray
        assert isinstance(expected, numpy.ndarray)

        testing.assert_array_equal(expected, actual)

    def test_module_not_callable(self):

        self.assertRaises(TypeError, fallback_mode.numpy)

        self.assertRaises(TypeError, fallback_mode.numpy.linalg)

    def test_numpy_scalars(self):

        assert fallback_mode.numpy.inf is numpy.inf

        assert fallback_mode.numpy.pi is numpy.pi

        # True, because 'is' checks for reference
        # fallback_mode.numpy.nan and numpy.nan both have same reference
        assert fallback_mode.numpy.nan is numpy.nan

    def test_cupy_specific_func(self):

        with self.assertRaises(AttributeError):
            func = fallback_mode.numpy.ElementwiseKernel  # NOQA

    def test_func_not_in_numpy(self):

        with self.assertRaises(AttributeError):
            func = fallback_mode.numpy.dummy  # NOQA

@testing.gpu
class TestNotifications(unittest.TestCase):

    def test_1_seterr_geterr(self):

        default = fallback_mode.geterr()
        assert default == 'warn'

        old = fallback_mode.seterr('ignore')
        current = fallback_mode.geterr()
        assert old == 'warn'
        assert current == 'ignore'

    def test_notification_ignore(self):

        import contextlib
        from io import StringIO

        saved_stdout = StringIO()
        with contextlib.redirect_stdout(saved_stdout):
            fallback_mode.seterr('ignore')
            res = fallback_mode.numpy.nanargmin([1, 2, 3])  # NOQA

        output = saved_stdout.getvalue().strip()
        assert output == ""

    def test_notification_print(self):

        import contextlib
        from io import StringIO

        saved_stdout = StringIO()
        with contextlib.redirect_stdout(saved_stdout):
            fallback_mode.seterr('print')
            res = fallback_mode.numpy.nanargmin([1, 2, 3])  # NOQA

        output = saved_stdout.getvalue().strip()
        assert output == ("Warning: 'nanargmin' method not in cupy, " +
                          "falling back to 'numpy.nanargmin'")

    def test_notification_warn(self):

        with self.assertWarns(fallback_mode.notification.FallbackWarning):
            fallback_mode.seterr('warn')
            res = fallback_mode.numpy.nanargmin([1, 2, 3])  # NOQA

    def test_notification_raise(self):

        with self.assertRaises(AttributeError):
            fallback_mode.seterr('raise')
            res = fallback_mode.numpy.nanargmin([1, 2, 3])  # NOQA

    def test_geterrcall(self):

        def f(func):
            pass

        fallback_mode.seterrcall(f)
        current = fallback_mode.geterrcall()

        assert current == f

    def test_notification_call(self):

        def custom_callback(func):
            print("'{}' fallbacked".format(func.__name__))

        import contextlib
        from io import StringIO

        saved_stdout = StringIO()
        with contextlib.redirect_stdout(saved_stdout):
            fallback_mode.seterrcall(custom_callback)
            fallback_mode.seterr('call')
            res = fallback_mode.numpy.nanargmin([1, 2, 3])  # NOQA

        output = saved_stdout.getvalue().strip()
        assert output == ("'nanargmin' fallbacked")

    def test_notification_log(self):

        class Log:
            def write(self, msg):
                print("LOG: {}".format(msg))

        import contextlib
        from io import StringIO

        saved_stdout = StringIO()
        with contextlib.redirect_stdout(saved_stdout):
            fallback_mode.seterrcall(Log())
            fallback_mode.seterr('log')
            res = fallback_mode.numpy.nanargmin([1, 2, 3])  # NOQA

        output = saved_stdout.getvalue().strip()
        assert output == ("LOG: 'nanargmin' method not in cupy, " +
                          "falling back to 'numpy.nanargmin'")

    def test_errstate(self):

        fallback_mode.seterr('print')
        before = fallback_mode.geterr()

        with fallback_mode.errstate('log'):
            inside = fallback_mode.geterr()
            assert inside == 'log'

        after = fallback_mode.geterr()
        assert before == after
<commit_msg>fix test flake error<commit_after>import unittest
import functools

import numpy

import cupy
from cupy import testing
from cupyx import fallback_mode

@testing.gpu
class TestFallbackMode(unittest.TestCase):

    def numpy_fallback_equal(name='xp'):
        """
        Decorator that checks fallback_mode results are equal to NumPy ones.

        Args:
            name(str): Argument name whose value is either
            ``numpy`` or ``cupy`` module.
        """
        def decorator(impl):
            @functools.wraps(impl)
            def test_func(self, *args, **kwargs):

                kwargs[name] = fallback_mode.numpy
                fallback_result = impl(self, *args, **kwargs)

                kwargs[name] = numpy
                numpy_result = impl(self, *args, **kwargs)

                if isinstance(numpy_result, numpy.ndarray):
                    # if numpy returns ndarray, cupy must return ndarray
                    assert isinstance(fallback_result, cupy.ndarray)
                    testing.assert_array_equal(numpy_result, fallback_result)

                elif isinstance(numpy_result, numpy.ScalarType):
                    # if numpy returns scalar
                    # cupy must return scalar or 0-dim array
                    if isinstance(fallback_result, numpy.ScalarType):
                        assert numpy_result == fallback_result

                    else:
                        # cupy 0-dim array
                        assert numpy_result == int(fallback_result)
                else:
                    raise NotImplementedError

            return test_func
        return decorator

    @numpy_fallback_equal()
    def test_argmin(self, xp):

        a = xp.array([
            [13, 5, 45, 23, 9],
            [-5, 41, 0, 22, 4],
            [2, 6, 43, 11, -1]
        ])

        return xp.argmin(a, axis=1)

    @numpy_fallback_equal()
    def test_argmin_zero_dim_array_vs_scalar(self, xp):

        a = xp.array([
            [13, 5, 45, 23, 9],
            [-5, 41, 0, 22, 4],
            [2, 6, 43, 11, -1]
        ])

        return xp.argmin(a)

    # cupy.argmin raises error if list passed, numpy does not
    def test_argmin_list(self):

        a = [
            [13, 5, 45, 23, 9],
            [-5, 41, 0, 22, 4],
            [2, 6, 43, 11, -1]
        ]

        with self.assertRaises(AttributeError):
            fallback_mode.numpy.argmin(a)

        assert numpy.argmin([1, 0, 3]) == 1

    # Non-existing function
    @numpy_fallback_equal()
    def test_array_equal(self, xp):

        a = xp.array([1, 2])
        b = xp.array([1, 2])

        return xp.array_equal(a, b)

    # Both cupy and numpy return 0-d array
    @numpy_fallback_equal()
    def test_convolve_zero_dim_array(self, xp):

        a = xp.array([1, 2, 3])
        b = xp.array([0, 1, 0.5])

        return xp.convolve(a, b, 'valid')

    def test_vectorize(self):

        def function(a, b):
            if a > b:
                return a - b
            return a + b

        actual = numpy.vectorize(function)([1, 2, 3, 4], 2)
        expected = fallback_mode.numpy.vectorize(function)([1, 2, 3, 4], 2)

        assert isinstance(actual, numpy.ndarray)

        # ([1,2,3,4], 2) are arguments to
        # numpy.vectorize(function), not numpy.vectorize
        # So, it returns numpy.ndarray
        assert isinstance(expected, numpy.ndarray)

        testing.assert_array_equal(expected, actual)

    def test_module_not_callable(self):

        self.assertRaises(TypeError, fallback_mode.numpy)

        self.assertRaises(TypeError, fallback_mode.numpy.linalg)

    def test_numpy_scalars(self):

        assert fallback_mode.numpy.inf is numpy.inf

        assert fallback_mode.numpy.pi is numpy.pi

        # True, because 'is' checks for reference
        # fallback_mode.numpy.nan and numpy.nan both have same reference
        assert fallback_mode.numpy.nan is numpy.nan

    def test_cupy_specific_func(self):

        with self.assertRaises(AttributeError):
            func = fallback_mode.numpy.ElementwiseKernel  # NOQA

    def test_func_not_in_numpy(self):

        with self.assertRaises(AttributeError):
            func = fallback_mode.numpy.dummy  # NOQA

@testing.gpu
class TestNotifications(unittest.TestCase):

    def test_1_seterr_geterr(self):

        default = fallback_mode.geterr()
        assert default == 'warn'

        old = fallback_mode.seterr('ignore')
        current = fallback_mode.geterr()
        assert old == 'warn'
        assert current == 'ignore'

    def test_notification_ignore(self):

        import contextlib
        from io import StringIO

        saved_stdout = StringIO()
        with contextlib.redirect_stdout(saved_stdout):
            fallback_mode.seterr('ignore')
            fallback_mode.numpy.nanargmin([1, 2, 3])

        output = saved_stdout.getvalue().strip()
        assert output == ""

    def test_notification_print(self):

        import contextlib
        from io import StringIO

        saved_stdout = StringIO()
        with contextlib.redirect_stdout(saved_stdout):
            fallback_mode.seterr('print')
            res = fallback_mode.numpy.nanargmin([1, 2, 3])  # NOQA

        output = saved_stdout.getvalue().strip()
        assert output == ("Warning: 'nanargmin' method not in cupy, " +
                          "falling back to 'numpy.nanargmin'")

    def test_notification_warn(self):

        with self.assertWarns(fallback_mode.notification.FallbackWarning):
            fallback_mode.seterr('warn')
            res = fallback_mode.numpy.nanargmin([1, 2, 3])  # NOQA

    def test_notification_raise(self):

        with self.assertRaises(AttributeError):
            fallback_mode.seterr('raise')
            res = fallback_mode.numpy.nanargmin([1, 2, 3])  # NOQA

    def test_geterrcall(self):

        def f(func):
            pass

        fallback_mode.seterrcall(f)
        current = fallback_mode.geterrcall()

        assert current == f

    def test_notification_call(self):

        def custom_callback(func):
            print("'{}' fallbacked".format(func.__name__))

        import contextlib
        from io import StringIO

        saved_stdout = StringIO()
        with contextlib.redirect_stdout(saved_stdout):
            fallback_mode.seterrcall(custom_callback)
            fallback_mode.seterr('call')
            res = fallback_mode.numpy.nanargmin([1, 2, 3])  # NOQA

        output = saved_stdout.getvalue().strip()
        assert output == ("'nanargmin' fallbacked")

    def test_notification_log(self):

        class Log:
            def write(self, msg):
                print("LOG: {}".format(msg))

        import contextlib
        from io import StringIO

        saved_stdout = StringIO()
        with contextlib.redirect_stdout(saved_stdout):
            fallback_mode.seterrcall(Log())
            fallback_mode.seterr('log')
            res = fallback_mode.numpy.nanargmin([1, 2, 3])  # NOQA

        output = saved_stdout.getvalue().strip()
        assert output == ("LOG: 'nanargmin' method not in cupy, " +
                          "falling back to 'numpy.nanargmin'")

    def test_errstate(self):

        fallback_mode.seterr('print')
        before = fallback_mode.geterr()

        with fallback_mode.errstate('log'):
            inside = fallback_mode.geterr()
            assert inside == 'log'

        after = fallback_mode.geterr()
        assert before == after
<|endoftext|>


<commit_before>#!/usr/bin/env python
import os
import sys

import attr
from gevent import monkey
from psycogreen.gevent import patch_psycopg

def main():
    patch_pickle()  # should happen before gevent patch, which imports pickle

    # important to apply gevent monkey patches before running any other code
    # applying this later can lead to inconsistencies and threading issues
    # but compressor doesn't like it
    # ('module' object has no attribute 'poll' which has to do with
    # gevent-patching subprocess)
    GEVENT_COMMANDS = (
        GeventCommand('run_gunicorn'),
        GeventCommand('run_sql'),
        GeventCommand('run_blob_migration'),
        GeventCommand('check_blob_logs'),
        GeventCommand('preindex_everything'),
        GeventCommand('migrate_multi'),
        GeventCommand('prime_views'),
        GeventCommand('ptop_preindex'),
        GeventCommand('sync_prepare_couchdb_multi'),
        GeventCommand('sync_couch_views'),
        GeventCommand('populate_form_date_modified'),
        GeventCommand('run_aggregation_query'),
        GeventCommand('send_pillow_retry_queue_through_pillows'),
        GeventCommand('run_all_management_command'),
        GeventCommand('copy_events_to_sql', http_adapter_pool_size=32)
    )
    _patch_gevent_if_required(sys.argv, GEVENT_COMMANDS)

    init_hq_python_path()
    run_patches()

    set_default_settings_path(sys.argv)
    set_nosetests_verbosity(sys.argv)
    from django.core.management import execute_from_command_line
    execute_from_command_line(sys.argv)

@attr.s
class GeventCommand(object):
    command = attr.ib()
    contains = attr.ib(default=None)
    http_adapter_pool_size = attr.ib(default=None)

def _patch_gevent_if_required(args, gevent_commands):
    if len(args) <= 1:
        return
    for gevent_command in gevent_commands:
        should_patch = args[1] == gevent_command.command
        if gevent_command.contains:
            should_patch = should_patch and gevent_command.contains in ' '.join(args)
        if should_patch:
            monkey.patch_all(subprocess=True)
            patch_psycopg()
            if gevent_command.http_adapter_pool_size:
                # requests must be imported after `patch_all()` is called
                import requests
                requests.adapters.DEFAULT_POOLSIZE = gevent_command.http_adapter_pool_size
            break

def init_hq_python_path():
    _set_source_root_parent('submodules')
    _set_source_root_parent('extensions')
    _set_source_root(os.path.join('corehq', 'ex-submodules'))
    _set_source_root(os.path.join('custom', '_legacy'))

def _set_source_root_parent(source_root_parent):
    """
    add everything under `source_root_parent` to the list of source roots
    e.g. if you call this with param 'submodules'
    and you have the file structure

    project/
        submodules/
            foo-src/
                foo/
            bar-src/
                bar/

    (where foo and bar are python modules)
    then foo and bar would become top-level importable

    """
    filedir = os.path.dirname(__file__)
    dir = os.path.join(filedir, source_root_parent)
    if not os.path.exists(dir):
        return

    submodules_list = os.listdir(dir)
    for d in submodules_list:
        if d == "__init__.py" or d == '.' or d == '..':
            continue
        sys.path.insert(1, os.path.join(filedir, source_root_parent, d))

    sys.path.append(os.path.join(filedir, source_root_parent))

def _set_source_root(source_root):
    filedir = os.path.dirname(__file__)
    sys.path.insert(1, os.path.join(filedir, source_root))

def run_patches():
    # workaround for https://github.com/smore-inc/tinys3/issues/33
    import mimetypes
    mimetypes.init()

    patch_pickle()
    patch_jsonfield()

    import django
    _setup_once.setup = django.setup
    django.setup = _setup_once

def patch_pickle():
    """Patch pickle to support protocol 5

    Remove after upgrading to Python 3.8+
    """
    import pickle
    if pickle.HIGHEST_PROTOCOL < 5:
        import pickle5
        sys.modules["pickle"] = pickle5

def patch_jsonfield():
    """Patch the ``to_python`` method of JSONField
    See https://github.com/bradjasper/django-jsonfield/pull/173 for more details
    """
    import json
    from django.core.exceptions import ValidationError
    from django.utils.translation import gettext_lazy as _
    from jsonfield import JSONField

    def to_python(self, value):
        if isinstance(value, str):
            try:
                return json.loads(value, **self.load_kwargs)
            except ValueError:
                raise ValidationError(_("Enter valid JSON"))
        return value

    JSONField.to_python = to_python

# HACK monkey-patch django setup to prevent second setup by django_nose
def _setup_once(*args, **kw):
    if not hasattr(_setup_once, "done"):
        _setup_once.done = True
        _setup_once.setup(*args, **kw)

def set_default_settings_path(argv):
    if len(argv) > 1 and argv[1] == 'test' or os.environ.get('CCHQ_TESTING') == '1':
        os.environ.setdefault('CCHQ_TESTING', '1')
        module = 'testsettings'
    else:
        module = 'settings'
    os.environ.setdefault("DJANGO_SETTINGS_MODULE", module)

def set_nosetests_verbosity(argv):
    """Increase nose output verbosity with -v... argument

        -v: print test names
        -vv: do not capture stdout
        -vvv: do not capture logging
        -vvvv: enable nose internal logging
    """
    import logging

    def set_verbosity(arg, i):
        args = []
        verbosity = sum(1 for c in arg if c == "v") + 1
        if len(arg) > verbosity:
            # preserve other single-letter arguments (ex: -xv)
            args.append("".join(c for c in arg if c != "v"))
        if verbosity > 2:
            args.append("--nocapture")
        if verbosity > 3:
            verbosity -= 1
            args.append("--nologcapture")
            logging.basicConfig(level=logging.NOTSET)
            logging.getLogger().info(
                "Adjust logging with testsettings._set_logging_levels")
        args.append("--nose-verbosity=%s" % verbosity)
        argv[i:i + 1] = args

    if len(argv) > 1 and argv[1] == 'test':
        for i, arg in reversed(list(enumerate(argv))):
            if arg[:1] == "-" and arg[1] != "-" and any(c == 'v' for c in arg):
                set_verbosity(arg, i)
                break

if __name__ == "__main__":
    main()
<commit_msg>Patch for Django < 3.1<commit_after>#!/usr/bin/env python
import os
import sys

import attr
from gevent import monkey
from psycogreen.gevent import patch_psycopg

def main():
    patch_pickle()  # should happen before gevent patch, which imports pickle

    # important to apply gevent monkey patches before running any other code
    # applying this later can lead to inconsistencies and threading issues
    # but compressor doesn't like it
    # ('module' object has no attribute 'poll' which has to do with
    # gevent-patching subprocess)
    GEVENT_COMMANDS = (
        GeventCommand('run_gunicorn'),
        GeventCommand('run_sql'),
        GeventCommand('run_blob_migration'),
        GeventCommand('check_blob_logs'),
        GeventCommand('preindex_everything'),
        GeventCommand('migrate_multi'),
        GeventCommand('prime_views'),
        GeventCommand('ptop_preindex'),
        GeventCommand('sync_prepare_couchdb_multi'),
        GeventCommand('sync_couch_views'),
        GeventCommand('populate_form_date_modified'),
        GeventCommand('run_aggregation_query'),
        GeventCommand('send_pillow_retry_queue_through_pillows'),
        GeventCommand('run_all_management_command'),
        GeventCommand('copy_events_to_sql', http_adapter_pool_size=32)
    )
    _patch_gevent_if_required(sys.argv, GEVENT_COMMANDS)

    init_hq_python_path()
    run_patches()

    set_default_settings_path(sys.argv)
    set_nosetests_verbosity(sys.argv)
    from django.core.management import execute_from_command_line
    execute_from_command_line(sys.argv)

@attr.s
class GeventCommand(object):
    command = attr.ib()
    contains = attr.ib(default=None)
    http_adapter_pool_size = attr.ib(default=None)

def _patch_gevent_if_required(args, gevent_commands):
    if len(args) <= 1:
        return
    for gevent_command in gevent_commands:
        should_patch = args[1] == gevent_command.command
        if gevent_command.contains:
            should_patch = should_patch and gevent_command.contains in ' '.join(args)
        if should_patch:
            monkey.patch_all(subprocess=True)
            patch_psycopg()
            if gevent_command.http_adapter_pool_size:
                # requests must be imported after `patch_all()` is called
                import requests
                requests.adapters.DEFAULT_POOLSIZE = gevent_command.http_adapter_pool_size
            break

def init_hq_python_path():
    _set_source_root_parent('submodules')
    _set_source_root_parent('extensions')
    _set_source_root(os.path.join('corehq', 'ex-submodules'))
    _set_source_root(os.path.join('custom', '_legacy'))

def _set_source_root_parent(source_root_parent):
    """
    add everything under `source_root_parent` to the list of source roots
    e.g. if you call this with param 'submodules'
    and you have the file structure

    project/
        submodules/
            foo-src/
                foo/
            bar-src/
                bar/

    (where foo and bar are python modules)
    then foo and bar would become top-level importable

    """
    filedir = os.path.dirname(__file__)
    dir = os.path.join(filedir, source_root_parent)
    if not os.path.exists(dir):
        return

    submodules_list = os.listdir(dir)
    for d in submodules_list:
        if d == "__init__.py" or d == '.' or d == '..':
            continue
        sys.path.insert(1, os.path.join(filedir, source_root_parent, d))

    sys.path.append(os.path.join(filedir, source_root_parent))

def _set_source_root(source_root):
    filedir = os.path.dirname(__file__)
    sys.path.insert(1, os.path.join(filedir, source_root))

def run_patches():
    # workaround for https://github.com/smore-inc/tinys3/issues/33
    import mimetypes
    mimetypes.init()

    patch_pickle()
    patch_jsonfield()

    import django
    _setup_once.setup = django.setup
    django.setup = _setup_once

def patch_pickle():
    """Patch pickle to support protocol 5

    Remove after upgrading to Python 3.8+
    """
    import pickle
    if pickle.HIGHEST_PROTOCOL < 5:
        import pickle5
        sys.modules["pickle"] = pickle5

def patch_jsonfield():
    """Patch the ``to_python`` method of JSONField
    See https://github.com/bradjasper/django-jsonfield/pull/173 for more details
    """
    import json
    from django.core.exceptions import ValidationError
    from django.utils.translation import gettext_lazy as _
    from jsonfield import JSONField

    def to_python(self, value):
        if isinstance(value, str):
            try:
                return json.loads(value, **self.load_kwargs)
            except ValueError:
                raise ValidationError(_("Enter valid JSON"))
        return value

    JSONField.to_python = to_python

    import django
    if django.VERSION < (3, 1):
        # TODO remove after upgrading to Django 3.2
        from django.contrib.postgres.fields.jsonb import JSONField
        import django.db.models
        django.db.models.JSONField = JSONField

# HACK monkey-patch django setup to prevent second setup by django_nose
def _setup_once(*args, **kw):
    if not hasattr(_setup_once, "done"):
        _setup_once.done = True
        _setup_once.setup(*args, **kw)

def set_default_settings_path(argv):
    if len(argv) > 1 and argv[1] == 'test' or os.environ.get('CCHQ_TESTING') == '1':
        os.environ.setdefault('CCHQ_TESTING', '1')
        module = 'testsettings'
    else:
        module = 'settings'
    os.environ.setdefault("DJANGO_SETTINGS_MODULE", module)

def set_nosetests_verbosity(argv):
    """Increase nose output verbosity with -v... argument

        -v: print test names
        -vv: do not capture stdout
        -vvv: do not capture logging
        -vvvv: enable nose internal logging
    """
    import logging

    def set_verbosity(arg, i):
        args = []
        verbosity = sum(1 for c in arg if c == "v") + 1
        if len(arg) > verbosity:
            # preserve other single-letter arguments (ex: -xv)
            args.append("".join(c for c in arg if c != "v"))
        if verbosity > 2:
            args.append("--nocapture")
        if verbosity > 3:
            verbosity -= 1
            args.append("--nologcapture")
            logging.basicConfig(level=logging.NOTSET)
            logging.getLogger().info(
                "Adjust logging with testsettings._set_logging_levels")
        args.append("--nose-verbosity=%s" % verbosity)
        argv[i:i + 1] = args

    if len(argv) > 1 and argv[1] == 'test':
        for i, arg in reversed(list(enumerate(argv))):
            if arg[:1] == "-" and arg[1] != "-" and any(c == 'v' for c in arg):
                set_verbosity(arg, i)
                break

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


<commit_before>import asyncio
import json
import logging
from collections import Mapping, Sequence

from aiohttp import client, ClientOSError
from yarl import URL

from . import base, StorageError

class RoStorage(base.AbstractStorageReadOnly):
    """ ReadOnly storage over http GET
    config:
        semaphore: int
        allow_hosts: list
        return_status: bool, method get returns tuple (CODE, VALUE)
        prefix: url prefix
        headers: Mapping or None
        template: url template
        format: [json|str|bytes], default json
    """
    async def init(self):
        self._prefix = self.config.get('prefix')
        self._template = self.config.get('template')
        if self._prefix:
            self._prefix = URL(self._prefix)
        self._semaphore = asyncio.Semaphore(
            self.config.get('semaphore', 20), loop=self.loop)
        self._allow_hosts = self.config.get('allow_hosts')
        self._format = self.config.get('format', 'json')
        self._return_status = self.config.get('return_status', False)

        headers = self.config.get('headers')
        self.session_params = {}
        if headers:
            self.session_params['headers'] = dict(headers)
        for param in ('conn_timeout', 'read_timeout'):
            if param in self.config:
                self.session_params[param] = self.config[param]
        self.reset_session()
        self.context.on_stop.append(self.stop)

    def reset_session(self, **kwargs):
        session = getattr(self, 'session', None)
        if session:
            session.close()
        if kwargs:
            kwargs = {**self.session_params, **kwargs}
        else:
            kwargs = self.session_params
        self.session = client.ClientSession(loop=self.loop, **kwargs)

    async def stop(self):
        self.session.close()

    def raw_key(self, key):
        if self._prefix:
            url = self._prefix / key
        elif self._template and isinstance(key, Mapping):
            url = URL(self._template.format_map(key))
        elif self._template and isinstance(key, Sequence):
            url = URL(self._template.format(*key))
        elif self._template:
            url = URL(self._template.format(key))
        elif isinstance(key, str):
            url = URL(key)
        else:
            url = key
        if self._allow_hosts and url.host not in self._allow_hosts:
            raise KeyError(key)
        return url

    async def _request(self, url, *, method='get', **kwargs):
        async with self._semaphore:
            coro = getattr(self.session, method)
            async with coro(url, **kwargs) as response:
                if self._format == 'json' and 'json' in response.content_type:
                    return response.status, await response.json()
                elif self._format == 'str':
                    return response.status, await response.text()
                else:
                    return response.status, await response.read()

    async def request(self, url, **kwargs):
        try:
            status, data = await self._request(url, **kwargs)
        except ClientOSError as e:
            raise StorageError('URL %s: %s' % (url, e)) from e

        if self._return_status:
            return status, data

        if status == 404:
            data = None
        elif status >= 400:
            raise StorageError('URL %s: %s' % (url, status))
        return data

    def get(self, key):
        url = self.raw_key(key)
        return self.request(url)

    async def copy(self, key_source, storage_dest, key_dest):
        """ Return True if data are copied
        * optimized for http->fs copy
        * not supported return_status
        """
        from aioworkers.storage.filesystem import FileSystemStorage
        if not isinstance(storage_dest, FileSystemStorage):
            return super().copy(key_source, storage_dest, key_dest)
        url = self.raw_key(key_source)
        logger = self.context.logger
        async with self._semaphore:
            async with self.session.get(url) as response:
                if response.status == 404:
                    return
                elif response.status >= 400:
                    if logger.getEffectiveLevel() == logging.DEBUG:
                        logger.debug(
                            'HttpStorage request to %s '
                            'returned code %s:\n%s' % (
                                url, response.status,
                                (await response.read()).decode()))
                    return
                async with storage_dest.raw_key(key_dest).open('wb') as f:
                    async for chunk in response.content.iter_any():
                        await f.write(chunk)
                    return True

class Storage(RoStorage, base.AbstractStorageWriteOnly):
    """ RW storage over http
    config:
        semaphore: int
        allow_hosts: list
        return_status: bool, method get returns tuple (CODE, VALUE)
        prefix: url prefix
        template: url template
        headers: Mapping or None
        format: [json|str|bytes], default json
        set: [post|put|patch], default post
        dumps: str, path in context to dumps
    """

    def set(self, key, value):
        url = self.raw_key(key)
        if self._format == 'json':
            if self.config.get('dumps'):
                dumps = self.context[self.config.dumps]
            else:
                dumps = json.dumps
            data = dumps(value)
            headers = {'content-type': 'application/json'}
        else:
            data = value
            headers = {}

        return self.request(
            url, method=self.config.get('set', 'post'),
            data=data, headers=headers)
<commit_msg>fix aiohttp 3.0.X compatibility http.Storage<commit_after>import asyncio
import json
import logging
from collections import Mapping, Sequence

from aiohttp import client, ClientOSError
from yarl import URL

from . import base, StorageError

class RoStorage(base.AbstractStorageReadOnly):
    """ ReadOnly storage over http GET
    config:
        semaphore: int
        allow_hosts: list
        return_status: bool, method get returns tuple (CODE, VALUE)
        prefix: url prefix
        headers: Mapping or None
        template: url template
        format: [json|str|bytes], default json
    """
    async def init(self):
        self._prefix = self.config.get('prefix')
        self._template = self.config.get('template')
        if self._prefix:
            self._prefix = URL(self._prefix)
        self._semaphore = asyncio.Semaphore(
            self.config.get('semaphore', 20), loop=self.loop)
        self._allow_hosts = self.config.get('allow_hosts')
        self._format = self.config.get('format', 'json')
        self._return_status = self.config.get('return_status', False)

        headers = self.config.get('headers')
        self.session_params = {}
        if headers:
            self.session_params['headers'] = dict(headers)
        for param in ('conn_timeout', 'read_timeout'):
            if param in self.config:
                self.session_params[param] = self.config[param]
        self.reset_session()
        self.context.on_stop.append(self.stop)

    def reset_session(self, **kwargs):
        session = getattr(self, 'session', None)
        if session:
            asyncio.ensure_future(session.close(), loop=self.loop)
        if kwargs:
            kwargs = {**self.session_params, **kwargs}
        else:
            kwargs = self.session_params
        self.session = client.ClientSession(loop=self.loop, **kwargs)

    async def stop(self):
        await self.session.close()

    def raw_key(self, key):
        if self._prefix:
            url = self._prefix / key
        elif self._template and isinstance(key, Mapping):
            url = URL(self._template.format_map(key))
        elif self._template and isinstance(key, Sequence):
            url = URL(self._template.format(*key))
        elif self._template:
            url = URL(self._template.format(key))
        elif isinstance(key, str):
            url = URL(key)
        else:
            url = key
        if self._allow_hosts and url.host not in self._allow_hosts:
            raise KeyError(key)
        return url

    async def _request(self, url, *, method='get', **kwargs):
        async with self._semaphore:
            coro = getattr(self.session, method)
            async with coro(url, **kwargs) as response:
                if self._format == 'json' and 'json' in response.content_type:
                    return response.status, await response.json()
                elif self._format == 'str':
                    return response.status, await response.text()
                else:
                    return response.status, await response.read()

    async def request(self, url, **kwargs):
        try:
            status, data = await self._request(url, **kwargs)
        except ClientOSError as e:
            raise StorageError('URL %s: %s' % (url, e)) from e

        if self._return_status:
            return status, data

        if status == 404:
            data = None
        elif status >= 400:
            raise StorageError('URL %s: %s' % (url, status))
        return data

    def get(self, key):
        url = self.raw_key(key)
        return self.request(url)

    async def copy(self, key_source, storage_dest, key_dest):
        """ Return True if data are copied
        * optimized for http->fs copy
        * not supported return_status
        """
        from aioworkers.storage.filesystem import FileSystemStorage
        if not isinstance(storage_dest, FileSystemStorage):
            return super().copy(key_source, storage_dest, key_dest)
        url = self.raw_key(key_source)
        logger = self.context.logger
        async with self._semaphore:
            async with self.session.get(url) as response:
                if response.status == 404:
                    return
                elif response.status >= 400:
                    if logger.getEffectiveLevel() == logging.DEBUG:
                        logger.debug(
                            'HttpStorage request to %s '
                            'returned code %s:\n%s' % (
                                url, response.status,
                                (await response.read()).decode()))
                    return
                async with storage_dest.raw_key(key_dest).open('wb') as f:
                    async for chunk in response.content.iter_any():
                        await f.write(chunk)
                    return True

class Storage(RoStorage, base.AbstractStorageWriteOnly):
    """ RW storage over http
    config:
        semaphore: int
        allow_hosts: list
        return_status: bool, method get returns tuple (CODE, VALUE)
        prefix: url prefix
        template: url template
        headers: Mapping or None
        format: [json|str|bytes], default json
        set: [post|put|patch], default post
        dumps: str, path in context to dumps
    """

    def set(self, key, value):
        url = self.raw_key(key)
        if self._format == 'json':
            if self.config.get('dumps'):
                dumps = self.context[self.config.dumps]
            else:
                dumps = json.dumps
            data = dumps(value)
            headers = {'content-type': 'application/json'}
        else:
            data = value
            headers = {}

        return self.request(
            url, method=self.config.get('set', 'post'),
            data=data, headers=headers)
<|endoftext|>


<commit_before>"""
Custom model fields to link to CMS content.
"""
import django
from django.utils import six
from django.core.exceptions import ValidationError
from django.db import models
from any_urlfield.models.values import AnyUrlValue
from any_urlfield.registry import UrlTypeRegistry
from any_urlfield.validators import ExtendedURLValidator

if django.VERSION < (1, 8):
    base_class = six.with_metaclass(models.SubfieldBase, models.CharField)
else:
    base_class = models.CharField

class AnyUrlField(base_class):
    """
    A CharField that can either refer to a CMS page ID, or external URL.

    .. figure:: /images/anyurlfield1.*
       :width: 363px
       :height: 74px
       :alt: AnyUrlField, with external URL input.

    .. figure:: /images/anyurlfield2.*
       :width: 290px
       :height: 76px
       :alt: AnyUrlField, with internal page input.

    By default, the ``AnyUrlField`` only supports linking to external pages.
    To add support for your own models (e.g. an ``Article`` model),
    include the following code in :file:`models.py`:

    .. code-block:: python

        from any_urlfield.models import AnyUrlField
        AnyUrlField.register_model(Article)

    Now, the ``AnyUrlField`` offers users a dropdown field to directly select an article.
    By default, it uses a :class:`django.forms.ModelChoiceField` field with a :class:`django.forms.Select` widget
    to render the field.  This can be customized using the ``form_field`` and ``widget`` parameters:

    .. code-block:: python

        from any_urlfield.models import AnyUrlField
        from any_urlfield.forms import SimpleRawIdWidget

        AnyUrlField.register_model(Article, widget=SimpleRawIdWidget(Article))

    Now, the ``Article`` model will be displayed as raw input field with a browse button.
    """
    _static_registry = UrlTypeRegistry()  # Also accessed by AnyUrlValue as internal field.

    def __init__(self, *args, **kwargs):
        if 'max_length' not in kwargs:
            kwargs['max_length'] = 300
        super(AnyUrlField, self).__init__(*args, **kwargs)

    @classmethod
    def register_model(cls, ModelClass, form_field=None, widget=None, title=None, prefix=None):
        """
        Register a model to use in the URL field.

        This function needs to be called once for every model
        that should be selectable in the URL field.

        :param ModelClass: The model to register.
        :param form_field: The form field class used to render the field. This can be a lambda for lazy evaluation.
        :param widget: The widget class, can be used instead of the form field.
        :param title: The title of the model, by default it uses the models ``verbose_name``.
        :param prefix: A custom prefix for the model in the serialized database format. By default it uses "appname.modelname".
        """
        cls._static_registry.register(ModelClass, form_field, widget, title, prefix)

    def formfield(self, **kwargs):
        # Associate formfield.
        # Import locally to avoid circular references.
        from any_urlfield.forms.fields import AnyUrlField as AnyUrlFormField
        kwargs['form_class'] = AnyUrlFormField
        kwargs['url_type_registry'] = self._static_registry
        if 'widget' in kwargs:
            del kwargs['widget']
        return super(AnyUrlField, self).formfield(**kwargs)

    def from_db_value(self, value, expression, connection, context):
        # As of Django 1.8, this method is used to cast DB values to python values.
        # The call to to_python() is not used anymore.
        if value is None:
            return None
        return AnyUrlValue.from_db_value(value, self._static_registry)

    def to_python(self, value):
        if isinstance(value, AnyUrlValue):
            return value

        # Convert the string value
        if value is None:
            return None

        return AnyUrlValue.from_db_value(value, self._static_registry)

    def get_prep_value(self, value):
        if isinstance(value, six.string_types):
            # Happens with south migration
            return value
        elif value is None:
            return None if self.null else ''
        else:
            # Convert back to string
            return value.to_db_value()

    def pre_save(self, model_instance, add):
        # Make sure that the SQL compiler in Django 1.6/1.7 doesn't get a AnyUrlValue,
        # but a regular 'str' object it can write to the database.
        value = super(AnyUrlField, self).pre_save(model_instance, add)
        if not value:
            return None
        else:
            return value.to_db_value()

    def value_to_string(self, obj):
        # For dumpdata and serialization
        value = self._get_val_from_obj(obj)
        return self.get_prep_value(value)

    def validate(self, value, model_instance):
        # Final validation of the field, before storing in the DB.
        super(AnyUrlField, self).validate(value, model_instance)
        if value:
            if value.type_prefix == 'http':
                validate_url = ExtendedURLValidator()
                validate_url(value.type_value)
            elif value.type_value:
                if not value.exists():
                    raise ValidationError(self.error_messages['invalid_choice'] % value.type_value)
<commit_msg>Replace _get_val_from_obj() with self.value_from_object() for Django 2.0<commit_after>"""
Custom model fields to link to CMS content.
"""
import django
from django.utils import six
from django.core.exceptions import ValidationError
from django.db import models
from any_urlfield.models.values import AnyUrlValue
from any_urlfield.registry import UrlTypeRegistry
from any_urlfield.validators import ExtendedURLValidator

if django.VERSION < (1, 8):
    base_class = six.with_metaclass(models.SubfieldBase, models.CharField)
else:
    base_class = models.CharField

class AnyUrlField(base_class):
    """
    A CharField that can either refer to a CMS page ID, or external URL.

    .. figure:: /images/anyurlfield1.*
       :width: 363px
       :height: 74px
       :alt: AnyUrlField, with external URL input.

    .. figure:: /images/anyurlfield2.*
       :width: 290px
       :height: 76px
       :alt: AnyUrlField, with internal page input.

    By default, the ``AnyUrlField`` only supports linking to external pages.
    To add support for your own models (e.g. an ``Article`` model),
    include the following code in :file:`models.py`:

    .. code-block:: python

        from any_urlfield.models import AnyUrlField
        AnyUrlField.register_model(Article)

    Now, the ``AnyUrlField`` offers users a dropdown field to directly select an article.
    By default, it uses a :class:`django.forms.ModelChoiceField` field with a :class:`django.forms.Select` widget
    to render the field.  This can be customized using the ``form_field`` and ``widget`` parameters:

    .. code-block:: python

        from any_urlfield.models import AnyUrlField
        from any_urlfield.forms import SimpleRawIdWidget

        AnyUrlField.register_model(Article, widget=SimpleRawIdWidget(Article))

    Now, the ``Article`` model will be displayed as raw input field with a browse button.
    """
    _static_registry = UrlTypeRegistry()  # Also accessed by AnyUrlValue as internal field.

    def __init__(self, *args, **kwargs):
        if 'max_length' not in kwargs:
            kwargs['max_length'] = 300
        super(AnyUrlField, self).__init__(*args, **kwargs)

    @classmethod
    def register_model(cls, ModelClass, form_field=None, widget=None, title=None, prefix=None):
        """
        Register a model to use in the URL field.

        This function needs to be called once for every model
        that should be selectable in the URL field.

        :param ModelClass: The model to register.
        :param form_field: The form field class used to render the field. This can be a lambda for lazy evaluation.
        :param widget: The widget class, can be used instead of the form field.
        :param title: The title of the model, by default it uses the models ``verbose_name``.
        :param prefix: A custom prefix for the model in the serialized database format. By default it uses "appname.modelname".
        """
        cls._static_registry.register(ModelClass, form_field, widget, title, prefix)

    def formfield(self, **kwargs):
        # Associate formfield.
        # Import locally to avoid circular references.
        from any_urlfield.forms.fields import AnyUrlField as AnyUrlFormField
        kwargs['form_class'] = AnyUrlFormField
        kwargs['url_type_registry'] = self._static_registry
        if 'widget' in kwargs:
            del kwargs['widget']
        return super(AnyUrlField, self).formfield(**kwargs)

    def from_db_value(self, value, expression, connection, context):
        # As of Django 1.8, this method is used to cast DB values to python values.
        # The call to to_python() is not used anymore.
        if value is None:
            return None
        return AnyUrlValue.from_db_value(value, self._static_registry)

    def to_python(self, value):
        if isinstance(value, AnyUrlValue):
            return value

        # Convert the string value
        if value is None:
            return None

        return AnyUrlValue.from_db_value(value, self._static_registry)

    def get_prep_value(self, value):
        if isinstance(value, six.string_types):
            # Happens with south migration
            return value
        elif value is None:
            return None if self.null else ''
        else:
            # Convert back to string
            return value.to_db_value()

    def pre_save(self, model_instance, add):
        # Make sure that the SQL compiler in Django 1.6/1.7 doesn't get a AnyUrlValue,
        # but a regular 'str' object it can write to the database.
        value = super(AnyUrlField, self).pre_save(model_instance, add)
        if not value:
            return None
        else:
            return value.to_db_value()

    def value_to_string(self, obj):
        # For dumpdata and serialization
        value = self.value_from_object(obj)
        return self.get_prep_value(value)

    def validate(self, value, model_instance):
        # Final validation of the field, before storing in the DB.
        super(AnyUrlField, self).validate(value, model_instance)
        if value:
            if value.type_prefix == 'http':
                validate_url = ExtendedURLValidator()
                validate_url(value.type_value)
            elif value.type_value:
                if not value.exists():
                    raise ValidationError(self.error_messages['invalid_choice'] % value.type_value)
<|endoftext|>


<commit_before>import datetime

import pytest
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

import Pegasus.db.schema as schema
from Pegasus.db.ensembles import EMError, Triggers, TriggerType

@pytest.fixture(scope="function")
def session():
    """
    Create in-memory sqlite database with tables setup and return a db session
    object.
    """
    engine = create_engine("sqlite://")

    # create all tables in the schema
    schema.Base.metadata.create_all(engine)

    session = sessionmaker(bind=engine)()

    # create an ensemble entry
    session.add(
        schema.Ensemble(
            name="test-ensemble",
            created=datetime.datetime.now(),
            updated=datetime.datetime.now(),
            state="ACTIVE",
            max_running=1,
            max_planning=1,
            username="test-user",
        )
    )

    session.commit()

    yield session

    # close session, db will be released
    session.close()

class TestTriggers:
    def test_get_trigger(self, session):
        # insert trigger
        t = schema.Trigger(
            _id=1,
            ensemble_id=1,
            name="test-trigger",
            state="STOPPED",
            workflow=r'{"script":"/wf.py", "args":["arg1"]}',
            args=r'{"timeout":100, "interval":20}',
            _type=TriggerType.CRON.value,
        )
        session.add(t)

        triggers = Triggers(session)

        expected = {
            "id": 1,
            "ensemble_id": 1,
            "name": "test-trigger",
            "state": "STOPPED",
            "workflow": {"script": "/wf.py", "args": ["arg1"]},
            "args": {"timeout": 100, "interval": 20},
            "type": "CRON",
        }

        # get trigger and convert to dict for comparison
        result = Triggers.get_object(triggers.get_trigger(1, "test-trigger"))

        assert expected == result

    def test_get_trigger_not_found(self, session):
        with pytest.raises(EMError) as e:
            Triggers(session).get_trigger(1, "test-trigger")

        assert "No such trigger: test-trigger" in str(e)
        assert e.value.status_code == 404

    def test_list_triggers(self, session):
        t1 = schema.Trigger(
            _id=1,
            ensemble_id=1,
            name="test-trigger1",
            state="READY",
            workflow=r'{"script":"/wf.py", "args":["arg1"]}',
            args=r'{"timeout":100, "interval":20}',
            _type=TriggerType.CRON.value,
        )
        session.add(t1)

        t2 = schema.Trigger(
            _id=2,
            ensemble_id=1,
            name="test-trigger2",
            state="READY",
            workflow=r'{"script":"/wf.py", "args":["arg1"]}',
            args=r'{"timeout":100, "interval":20}',
            _type=TriggerType.CRON.value,
        )
        session.add(t2)
        session.commit()

        triggers = Triggers(session)
        result = triggers.list_triggers()
        assert len(result) == 2

    def test_list_triggers_by_ensemble(self, session):
        # add another ensemble to the ensemble table
        session.add(
            schema.Ensemble(
                id=2,
                name="test-ensemble2",
                created=datetime.datetime.now(),
                updated=datetime.datetime.now(),
                state="ACTIVE",
                max_running=1,
                max_planning=1,
                username="test-user",
            )
        )

        # add a trigger assigned to test-ensemble2
        t = schema.Trigger(
            _id=1,
            ensemble_id=2,
            name="test-trigger1",
            state="READY",
            workflow=r'{"script":"/wf.py", "args":["arg1"]}',
            args=r'{"timeout":100, "interval":20}',
            _type=TriggerType.CRON.value,
        )
        session.add(t)
        session.commit()

        triggers = Triggers(session)
        result = triggers.list_triggers_by_ensemble(
            username="test-user", ensemble="test-ensemble2"
        )

        assert len(result) == 1
        assert Triggers.get_object(result[0]) == {
            "id": 1,
            "ensemble_id": 2,
            "name": "test-trigger1",
            "state": "READY",
            "workflow": {"script": "/wf.py", "args": ["arg1"]},
            "args": {"timeout": 100, "interval": 20},
            "type": "CRON",
        }

        result = triggers.list_triggers_by_ensemble(
            username="test-user", ensemble="doesntexist"
        )
        assert len(result) == 0

    def test_insert_trigger(self, session):
        print(session.query(schema.Ensemble).all())
        triggers = Triggers(session)
        triggers.insert_trigger(
            ensemble_id=1,
            trigger="test-trigger",
            trigger_type=TriggerType.CRON.value,
            workflow_script="/wf.py",
            workflow_args=["arg1"],
            interval=10,
            timeout=20,
        )

        expected = {
            "id": 1,
            "ensemble_id": 1,
            "name": "test-trigger",
            "state": "READY",
            "workflow": {"script": "/wf.py", "args": ["arg1"]},
            "args": {"timeout": 20, "interval": 10},
            "type": "CRON",
        }

        result = Triggers.get_object(
            session.query(schema.Trigger)
            .filter_by(ensemble_id=1, name="test-trigger")
            .one()
        )

        assert expected == result

    def test_update_state(self, session):
        # insert trigger
        t = schema.Trigger(
            _id=1,
            ensemble_id=1,
            name="test-trigger",
            state="READY",
            workflow=r'{"script":"/wf.py", "args":["arg1"]}',
            args=r'{"timeout":100, "interval":20}',
            _type=TriggerType.CRON.value,
        )
        session.add(t)

        triggers = Triggers(session)
        triggers.update_state(ensemble_id=1, trigger_id=1, new_state="RUNNING")

        expected_state = "RUNNING"
        result = session.query(schema.Trigger).filter_by(_id=1).one().state

        assert expected_state == result

    def test_delete_trigger(self, session):
        # insert trigger
        t = schema.Trigger(
            _id=1,
            ensemble_id=1,
            name="test-trigger",
            state="READY",
            workflow=r'{"script":"/wf.py", "args":["arg1"]}',
            args=r'{"timeout":100, "interval":20}',
            _type=TriggerType.CRON.value,
        )
        session.add(t)
        assert len(session.query(schema.Trigger).all()) == 1

        triggers = Triggers(session)

        # delete trigger
        triggers.delete_trigger(ensemble_id=1, trigger="test-trigger")

        assert len(session.query(schema.Trigger).all()) == 0

    def test_get_object(self, session):
        t = schema.Trigger(
            _id=1,
            ensemble_id=1,
            name="test-trigger",
            state="READY",
            workflow=r'{"script":"/wf.py", "args":["arg1"]}',
            args=r'{"timeout":100, "interval":20}',
            _type=TriggerType.CRON.value,
        )

        expected = {
            "id": 1,
            "ensemble_id": 1,
            "name": "test-trigger",
            "state": "READY",
            "workflow": {"script": "/wf.py", "args": ["arg1"]},
            "args": {"timeout": 100, "interval": 20},
            "type": "CRON",
        }

        result = Triggers.get_object(t)

        assert expected == result
<commit_msg>fix in test_list_triggers_by_ensemble<commit_after>import datetime

import pytest
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

import Pegasus.db.schema as schema
from Pegasus.db.ensembles import EMError, Triggers, TriggerType

@pytest.fixture(scope="function")
def session():
    """
    Create in-memory sqlite database with tables setup and return a db session
    object.
    """
    engine = create_engine("sqlite://")

    # create all tables in the schema
    schema.Base.metadata.create_all(engine)

    session = sessionmaker(bind=engine)()

    # create an ensemble entry
    session.add(
        schema.Ensemble(
            name="test-ensemble",
            created=datetime.datetime.now(),
            updated=datetime.datetime.now(),
            state="ACTIVE",
            max_running=1,
            max_planning=1,
            username="test-user",
        )
    )

    session.commit()

    yield session

    # close session, db will be released
    session.close()

class TestTriggers:
    def test_get_trigger(self, session):
        # insert trigger
        t = schema.Trigger(
            _id=1,
            ensemble_id=1,
            name="test-trigger",
            state="STOPPED",
            workflow=r'{"script":"/wf.py", "args":["arg1"]}',
            args=r'{"timeout":100, "interval":20}',
            _type=TriggerType.CRON.value,
        )
        session.add(t)

        triggers = Triggers(session)

        expected = {
            "id": 1,
            "ensemble_id": 1,
            "name": "test-trigger",
            "state": "STOPPED",
            "workflow": {"script": "/wf.py", "args": ["arg1"]},
            "args": {"timeout": 100, "interval": 20},
            "type": "CRON",
        }

        # get trigger and convert to dict for comparison
        result = Triggers.get_object(triggers.get_trigger(1, "test-trigger"))

        assert expected == result

    def test_get_trigger_not_found(self, session):
        with pytest.raises(EMError) as e:
            Triggers(session).get_trigger(1, "test-trigger")

        assert "No such trigger: test-trigger" in str(e)
        assert e.value.status_code == 404

    def test_list_triggers(self, session):
        t1 = schema.Trigger(
            _id=1,
            ensemble_id=1,
            name="test-trigger1",
            state="READY",
            workflow=r'{"script":"/wf.py", "args":["arg1"]}',
            args=r'{"timeout":100, "interval":20}',
            _type=TriggerType.CRON.value,
        )
        session.add(t1)

        t2 = schema.Trigger(
            _id=2,
            ensemble_id=1,
            name="test-trigger2",
            state="READY",
            workflow=r'{"script":"/wf.py", "args":["arg1"]}',
            args=r'{"timeout":100, "interval":20}',
            _type=TriggerType.CRON.value,
        )
        session.add(t2)
        session.commit()

        triggers = Triggers(session)
        result = triggers.list_triggers()
        assert len(result) == 2

    def test_list_triggers_by_ensemble(self, session):
        # add another ensemble to the ensemble table
        session.add(
            schema.Ensemble(
                id=2,
                name="test-ensemble2",
                created=datetime.datetime.now(),
                updated=datetime.datetime.now(),
                state="ACTIVE",
                max_running=1,
                max_planning=1,
                username="test-user",
            )
        )
        session.commit()

        # add a trigger assigned to test-ensemble2
        t = schema.Trigger(
            _id=1,
            ensemble_id=2,
            name="test-trigger1",
            state="READY",
            workflow=r'{"script":"/wf.py", "args":["arg1"]}',
            args=r'{"timeout":100, "interval":20}',
            _type=TriggerType.CRON.value,
        )
        session.add(t)
        session.commit()

        triggers = Triggers(session)
        result = triggers.list_triggers_by_ensemble(
            username="test-user", ensemble="test-ensemble2"
        )

        assert len(result) == 1
        assert Triggers.get_object(result[0]) == {
            "id": 1,
            "ensemble_id": 2,
            "name": "test-trigger1",
            "state": "READY",
            "workflow": {"script": "/wf.py", "args": ["arg1"]},
            "args": {"timeout": 100, "interval": 20},
            "type": "CRON",
        }

        result = triggers.list_triggers_by_ensemble(
            username="test-user", ensemble="doesntexist"
        )
        assert len(result) == 0

    def test_insert_trigger(self, session):
        print(session.query(schema.Ensemble).all())
        triggers = Triggers(session)
        triggers.insert_trigger(
            ensemble_id=1,
            trigger="test-trigger",
            trigger_type=TriggerType.CRON.value,
            workflow_script="/wf.py",
            workflow_args=["arg1"],
            interval=10,
            timeout=20,
        )

        expected = {
            "id": 1,
            "ensemble_id": 1,
            "name": "test-trigger",
            "state": "READY",
            "workflow": {"script": "/wf.py", "args": ["arg1"]},
            "args": {"timeout": 20, "interval": 10},
            "type": "CRON",
        }

        result = Triggers.get_object(
            session.query(schema.Trigger)
            .filter_by(ensemble_id=1, name="test-trigger")
            .one()
        )

        assert expected == result

    def test_update_state(self, session):
        # insert trigger
        t = schema.Trigger(
            _id=1,
            ensemble_id=1,
            name="test-trigger",
            state="READY",
            workflow=r'{"script":"/wf.py", "args":["arg1"]}',
            args=r'{"timeout":100, "interval":20}',
            _type=TriggerType.CRON.value,
        )
        session.add(t)

        triggers = Triggers(session)
        triggers.update_state(ensemble_id=1, trigger_id=1, new_state="RUNNING")

        expected_state = "RUNNING"
        result = session.query(schema.Trigger).filter_by(_id=1).one().state

        assert expected_state == result

    def test_delete_trigger(self, session):
        # insert trigger
        t = schema.Trigger(
            _id=1,
            ensemble_id=1,
            name="test-trigger",
            state="READY",
            workflow=r'{"script":"/wf.py", "args":["arg1"]}',
            args=r'{"timeout":100, "interval":20}',
            _type=TriggerType.CRON.value,
        )
        session.add(t)
        assert len(session.query(schema.Trigger).all()) == 1

        triggers = Triggers(session)

        # delete trigger
        triggers.delete_trigger(ensemble_id=1, trigger="test-trigger")

        assert len(session.query(schema.Trigger).all()) == 0

    def test_get_object(self, session):
        t = schema.Trigger(
            _id=1,
            ensemble_id=1,
            name="test-trigger",
            state="READY",
            workflow=r'{"script":"/wf.py", "args":["arg1"]}',
            args=r'{"timeout":100, "interval":20}',
            _type=TriggerType.CRON.value,
        )

        expected = {
            "id": 1,
            "ensemble_id": 1,
            "name": "test-trigger",
            "state": "READY",
            "workflow": {"script": "/wf.py", "args": ["arg1"]},
            "args": {"timeout": 100, "interval": 20},
            "type": "CRON",
        }

        result = Triggers.get_object(t)

        assert expected == result
<|endoftext|>