<commit_before>import functools
from corehq.apps.app_manager import id_strings
from dimagi.utils.decorators.memoized import memoized
from corehq.apps.app_manager.util import is_sort_only_column
import langcodes
import commcare_translations
from corehq.apps.app_manager.templatetags.xforms_extras import clean_trans

def non_empty_only(dct):
    return dict([(key, value) for key, value in dct.items() if value])

def _create_custom_app_strings(app, lang, for_default=False):

    def trans(d):
        return clean_trans(d, langs)

    def maybe_add_index(text):
        if app.build_version >= '2.8':
            numeric_nav_on = app.profile.get('properties', {}).get('cc-entry-mode') == 'cc-entry-review'
            if app.profile.get('features', {}).get('sense') == 'true' or numeric_nav_on:
                text = "${0} %s" % (text,) if not (text and text[0].isdigit()) else text
        return text

    langs = [lang] + app.langs
    yield id_strings.homescreen_title(), app.name
    yield id_strings.app_display_name(), app.name

    yield 'cchq.case', "Case"
    yield 'cchq.referral', "Referral"

    # include language code names
    for lc in app.langs:
        name = langcodes.get_name(lc) or lc
        if name:
            yield lc, name

    for module in app.get_modules():
        for detail_type, detail, _ in module.get_details():
            if detail_type.startswith('case'):
                label = trans(module.case_label)
            elif detail_type.startswith('referral'):
                label = trans(module.referral_label)
            elif detail_type in ('product_short', 'product_long'):
                label = ''
            else:
                label = None
            if label is not None:
                yield id_strings.detail_title_locale(module, detail_type), label

            for column in detail.get_columns():
                if not is_sort_only_column(column):
                    yield id_strings.detail_column_header_locale(module, detail_type, column), trans(column.header)

                if column.format in ('enum', 'enum-image'):
                    for item in column.enum:
                        yield id_strings.detail_column_enum_variable(module, detail_type, column, item.key_as_variable), trans(item.value)
                elif column.format == "graph":
                    for index, item in enumerate(column.graph_configuration.annotations):
                        yield id_strings.graph_annotation(module, detail_type, column, index), trans(item.values)
                    for property, values in column.graph_configuration.locale_specific_config.iteritems():
                        yield id_strings.graph_configuration(module, detail_type, column, property), trans(values)

            for tab in detail.get_tabs():
                yield id_strings.detail_tab_title_locale(module, detail_type, tab), trans(tab.header)

        yield id_strings.module_locale(module), maybe_add_index(trans(module.name))

        icon = module.icon_app_string(lang, for_default=for_default)
        audio = module.audio_app_string(lang, for_default=for_default)
        if icon:
            yield id_strings.module_icon_locale(module), icon
        if audio:
            yield id_strings.module_audio_locale(module), audio

        if hasattr(module, 'report_configs'):
            yield id_strings.report_menu(), 'Reports'
            yield id_strings.report_name_header(), 'Report Name'
            yield id_strings.report_description_header(), 'Report Description'
            yield id_strings.report_last_sync(), 'Last Sync'
            for config in module.report_configs:
                yield id_strings.report_command(config.uuid), trans(config.header)
                yield id_strings.report_name(config.uuid), config.report.title
                for column in config.report.report_columns:
                    yield (
                        id_strings.report_column_header(config.uuid, column.column_id),
                        column.get_header(lang)
                    )

        if hasattr(module, 'case_list'):
            if module.case_list.show:
                yield id_strings.case_list_locale(module), trans(module.case_list.label) or "Case List"
                icon = module.case_list.icon_app_string(lang, for_default=for_default)
                audio = module.case_list.audio_app_string(lang, for_default=for_default)
                if icon:
                    yield id_strings.case_list_icon_locale(module), icon
                if audio:
                    yield id_strings.case_list_audio_locale(module), audio

        if hasattr(module, 'referral_list'):
            if module.referral_list.show:
                yield id_strings.referral_list_locale(module), trans(module.referral_list.label)
        for form in module.get_forms():
            form_name = trans(form.name) + ('${0}' if form.show_count else '')
            yield id_strings.form_locale(form), maybe_add_index(form_name)

            icon = form.icon_app_string(lang, for_default=for_default)
            audio = form.audio_app_string(lang, for_default=for_default)
            if icon:
                yield id_strings.form_icon_locale(form), icon
            if audio:
                yield id_strings.form_audio_locale(form), audio

        if hasattr(module, 'case_list_form') and module.case_list_form.form_id:
            yield (
                id_strings.case_list_form_locale(module),
                trans(module.case_list_form.label) or "Create a new Case"
            )
            icon = module.case_list_form.icon_app_string(lang, for_default=for_default)
            audio = module.case_list_form.audio_app_string(lang, for_default=for_default)
            if icon:
                yield id_strings.case_list_form_icon_locale(module), icon
            if audio:
                yield id_strings.case_list_form_audio_locale(module), audio

class AppStringsBase(object):

    def __init__(self, load_translations):
        self._load_translations = load_translations

    @memoized
    def get_default_translations(self, lang):
        return self._load_translations(lang)

    def create_custom_app_strings(self, app, lang, for_default=False):
        custom = dict(_create_custom_app_strings(app, lang, for_default=for_default))
        if not for_default:
            custom = non_empty_only(custom)
        return custom

    def create_app_strings(self, app, lang, for_default=False):
        messages = {}
        for part in self.app_strings_parts(app, lang, for_default=for_default):
            messages.update(part)
        return commcare_translations.dumps(messages).encode('utf-8')

    def app_strings_parts(self, app, lang, for_default=False):
        raise NotImplementedError()

    def create_default_app_strings(self, app):
        messages = {}

        for lc in reversed(app.langs):
            if lc == "default":
                continue
            new_messages = commcare_translations.loads(
                self.create_app_strings(app, lc, for_default=True)
            )

            for key, val in new_messages.items():
                # do not overwrite a real trans with a blank trans
                if not (val == '' and key in messages):
                    messages[key] = val

        if 'case_sharing.exactly_one_group' not in messages:
            messages['case_sharing.exactly_one_group'] = \
                (u'The case sharing settings for your user are incorrect. '
                 u'This user must be in exactly one case sharing group. Please contact your supervisor.')

        if 'case_autoload.fixture.exactly_one_fixture' not in messages:
            messages['case_autoload.fixture.exactly_one_fixture'] = \
                (u'The lookup table settings for your user are incorrect. '
                    u'This user must have access to exactly one lookup table row for the table: ${0}')

        from corehq.apps.app_manager.models import (
            AUTO_SELECT_CASE, AUTO_SELECT_FIXTURE, AUTO_SELECT_USER,
            AUTO_SELECT_LOCATION, AUTO_SELECT_USERCASE, AUTO_SELECT_RAW
        )

        mode_text = {
            AUTO_SELECT_FIXTURE: u'lookup table field',
            AUTO_SELECT_USER: u'user data key',
            AUTO_SELECT_CASE: u'case index',
            AUTO_SELECT_USERCASE: u'user case',
            AUTO_SELECT_RAW: u'custom xpath expression',
        }

        for mode, text in mode_text.items():
            key = 'case_autoload.{0}.property_missing'.format(mode)
            if key not in messages:
                messages[key] = (u'The {} specified for case auto-selecting '
                                 u'could not be found: ${{0}}').format(text)
            key = 'case_autoload.{0}.case_missing'.format(mode)
            if key not in messages:
                messages[key] = u'Unable to find case referenced by auto-select case ID.'

        key = 'case_autoload.{0}.property_missing'.format(AUTO_SELECT_LOCATION)
        messages[key] = (u"This form requires access to the user's location, "
                         "but none was found.")
        key = 'case_autoload.{0}.case_missing'.format(AUTO_SELECT_LOCATION)
        messages[key] = (u"This form requires the user's location to be "
                         "marked as 'Tracks Stock'.")

        return commcare_translations.dumps(messages).encode('utf-8')

class DumpKnownAppStrings(AppStringsBase):

    def app_strings_parts(self, app, lang, for_default=False):
        yield self.create_custom_app_strings(app, lang, for_default=for_default)
        yield self.get_default_translations(lang)
        yield non_empty_only(app.translations.get(lang, {}))

class SimpleAppStrings(AppStringsBase):

    def app_strings_parts(self, app, lang, for_default=False):
        yield self.create_custom_app_strings(app, lang, for_default=for_default)
        if not for_default:
            yield non_empty_only(app.translations.get(lang, {}))

class SelectKnownAppStrings(AppStringsBase):

    def get_app_translation_keys(self, app):
        return set.union(set(), *(
            set(t.keys()) for t in app.translations.values()
        ))

    def app_strings_parts(self, app, lang, for_default=False):
        yield self.create_custom_app_strings(app, lang,
                                             for_default=for_default)
        cc_trans = self.get_default_translations(lang)
        yield dict((key, cc_trans[key])
                   for key in self.get_app_translation_keys(app)
                   if key in cc_trans)
        yield non_empty_only(app.translations.get(lang, {}))

CHOICES = {
    'dump-known': DumpKnownAppStrings(commcare_translations.load_translations),
    'simple': SimpleAppStrings(None),
    'select-known': SelectKnownAppStrings(
        functools.partial(commcare_translations.load_translations, version=2)
    ),
}
<commit_msg>use report name given in app manager<commit_after>import functools
from corehq.apps.app_manager import id_strings
from dimagi.utils.decorators.memoized import memoized
from corehq.apps.app_manager.util import is_sort_only_column
import langcodes
import commcare_translations
from corehq.apps.app_manager.templatetags.xforms_extras import clean_trans

def non_empty_only(dct):
    return dict([(key, value) for key, value in dct.items() if value])

def _create_custom_app_strings(app, lang, for_default=False):

    def trans(d):
        return clean_trans(d, langs)

    def maybe_add_index(text):
        if app.build_version >= '2.8':
            numeric_nav_on = app.profile.get('properties', {}).get('cc-entry-mode') == 'cc-entry-review'
            if app.profile.get('features', {}).get('sense') == 'true' or numeric_nav_on:
                text = "${0} %s" % (text,) if not (text and text[0].isdigit()) else text
        return text

    langs = [lang] + app.langs
    yield id_strings.homescreen_title(), app.name
    yield id_strings.app_display_name(), app.name

    yield 'cchq.case', "Case"
    yield 'cchq.referral', "Referral"

    # include language code names
    for lc in app.langs:
        name = langcodes.get_name(lc) or lc
        if name:
            yield lc, name

    for module in app.get_modules():
        for detail_type, detail, _ in module.get_details():
            if detail_type.startswith('case'):
                label = trans(module.case_label)
            elif detail_type.startswith('referral'):
                label = trans(module.referral_label)
            elif detail_type in ('product_short', 'product_long'):
                label = ''
            else:
                label = None
            if label is not None:
                yield id_strings.detail_title_locale(module, detail_type), label

            for column in detail.get_columns():
                if not is_sort_only_column(column):
                    yield id_strings.detail_column_header_locale(module, detail_type, column), trans(column.header)

                if column.format in ('enum', 'enum-image'):
                    for item in column.enum:
                        yield id_strings.detail_column_enum_variable(module, detail_type, column, item.key_as_variable), trans(item.value)
                elif column.format == "graph":
                    for index, item in enumerate(column.graph_configuration.annotations):
                        yield id_strings.graph_annotation(module, detail_type, column, index), trans(item.values)
                    for property, values in column.graph_configuration.locale_specific_config.iteritems():
                        yield id_strings.graph_configuration(module, detail_type, column, property), trans(values)

            for tab in detail.get_tabs():
                yield id_strings.detail_tab_title_locale(module, detail_type, tab), trans(tab.header)

        yield id_strings.module_locale(module), maybe_add_index(trans(module.name))

        icon = module.icon_app_string(lang, for_default=for_default)
        audio = module.audio_app_string(lang, for_default=for_default)
        if icon:
            yield id_strings.module_icon_locale(module), icon
        if audio:
            yield id_strings.module_audio_locale(module), audio

        if hasattr(module, 'report_configs'):
            yield id_strings.report_menu(), 'Reports'
            yield id_strings.report_name_header(), 'Report Name'
            yield id_strings.report_description_header(), 'Report Description'
            yield id_strings.report_last_sync(), 'Last Sync'
            for config in module.report_configs:
                yield id_strings.report_command(config.uuid), trans(config.header)
                yield id_strings.report_name(config.uuid), trans(config.header)
                for column in config.report.report_columns:
                    yield (
                        id_strings.report_column_header(config.uuid, column.column_id),
                        column.get_header(lang)
                    )

        if hasattr(module, 'case_list'):
            if module.case_list.show:
                yield id_strings.case_list_locale(module), trans(module.case_list.label) or "Case List"
                icon = module.case_list.icon_app_string(lang, for_default=for_default)
                audio = module.case_list.audio_app_string(lang, for_default=for_default)
                if icon:
                    yield id_strings.case_list_icon_locale(module), icon
                if audio:
                    yield id_strings.case_list_audio_locale(module), audio

        if hasattr(module, 'referral_list'):
            if module.referral_list.show:
                yield id_strings.referral_list_locale(module), trans(module.referral_list.label)
        for form in module.get_forms():
            form_name = trans(form.name) + ('${0}' if form.show_count else '')
            yield id_strings.form_locale(form), maybe_add_index(form_name)

            icon = form.icon_app_string(lang, for_default=for_default)
            audio = form.audio_app_string(lang, for_default=for_default)
            if icon:
                yield id_strings.form_icon_locale(form), icon
            if audio:
                yield id_strings.form_audio_locale(form), audio

        if hasattr(module, 'case_list_form') and module.case_list_form.form_id:
            yield (
                id_strings.case_list_form_locale(module),
                trans(module.case_list_form.label) or "Create a new Case"
            )
            icon = module.case_list_form.icon_app_string(lang, for_default=for_default)
            audio = module.case_list_form.audio_app_string(lang, for_default=for_default)
            if icon:
                yield id_strings.case_list_form_icon_locale(module), icon
            if audio:
                yield id_strings.case_list_form_audio_locale(module), audio

class AppStringsBase(object):

    def __init__(self, load_translations):
        self._load_translations = load_translations

    @memoized
    def get_default_translations(self, lang):
        return self._load_translations(lang)

    def create_custom_app_strings(self, app, lang, for_default=False):
        custom = dict(_create_custom_app_strings(app, lang, for_default=for_default))
        if not for_default:
            custom = non_empty_only(custom)
        return custom

    def create_app_strings(self, app, lang, for_default=False):
        messages = {}
        for part in self.app_strings_parts(app, lang, for_default=for_default):
            messages.update(part)
        return commcare_translations.dumps(messages).encode('utf-8')

    def app_strings_parts(self, app, lang, for_default=False):
        raise NotImplementedError()

    def create_default_app_strings(self, app):
        messages = {}

        for lc in reversed(app.langs):
            if lc == "default":
                continue
            new_messages = commcare_translations.loads(
                self.create_app_strings(app, lc, for_default=True)
            )

            for key, val in new_messages.items():
                # do not overwrite a real trans with a blank trans
                if not (val == '' and key in messages):
                    messages[key] = val

        if 'case_sharing.exactly_one_group' not in messages:
            messages['case_sharing.exactly_one_group'] = \
                (u'The case sharing settings for your user are incorrect. '
                 u'This user must be in exactly one case sharing group. Please contact your supervisor.')

        if 'case_autoload.fixture.exactly_one_fixture' not in messages:
            messages['case_autoload.fixture.exactly_one_fixture'] = \
                (u'The lookup table settings for your user are incorrect. '
                    u'This user must have access to exactly one lookup table row for the table: ${0}')

        from corehq.apps.app_manager.models import (
            AUTO_SELECT_CASE, AUTO_SELECT_FIXTURE, AUTO_SELECT_USER,
            AUTO_SELECT_LOCATION, AUTO_SELECT_USERCASE, AUTO_SELECT_RAW
        )

        mode_text = {
            AUTO_SELECT_FIXTURE: u'lookup table field',
            AUTO_SELECT_USER: u'user data key',
            AUTO_SELECT_CASE: u'case index',
            AUTO_SELECT_USERCASE: u'user case',
            AUTO_SELECT_RAW: u'custom xpath expression',
        }

        for mode, text in mode_text.items():
            key = 'case_autoload.{0}.property_missing'.format(mode)
            if key not in messages:
                messages[key] = (u'The {} specified for case auto-selecting '
                                 u'could not be found: ${{0}}').format(text)
            key = 'case_autoload.{0}.case_missing'.format(mode)
            if key not in messages:
                messages[key] = u'Unable to find case referenced by auto-select case ID.'

        key = 'case_autoload.{0}.property_missing'.format(AUTO_SELECT_LOCATION)
        messages[key] = (u"This form requires access to the user's location, "
                         "but none was found.")
        key = 'case_autoload.{0}.case_missing'.format(AUTO_SELECT_LOCATION)
        messages[key] = (u"This form requires the user's location to be "
                         "marked as 'Tracks Stock'.")

        return commcare_translations.dumps(messages).encode('utf-8')

class DumpKnownAppStrings(AppStringsBase):

    def app_strings_parts(self, app, lang, for_default=False):
        yield self.create_custom_app_strings(app, lang, for_default=for_default)
        yield self.get_default_translations(lang)
        yield non_empty_only(app.translations.get(lang, {}))

class SimpleAppStrings(AppStringsBase):

    def app_strings_parts(self, app, lang, for_default=False):
        yield self.create_custom_app_strings(app, lang, for_default=for_default)
        if not for_default:
            yield non_empty_only(app.translations.get(lang, {}))

class SelectKnownAppStrings(AppStringsBase):

    def get_app_translation_keys(self, app):
        return set.union(set(), *(
            set(t.keys()) for t in app.translations.values()
        ))

    def app_strings_parts(self, app, lang, for_default=False):
        yield self.create_custom_app_strings(app, lang,
                                             for_default=for_default)
        cc_trans = self.get_default_translations(lang)
        yield dict((key, cc_trans[key])
                   for key in self.get_app_translation_keys(app)
                   if key in cc_trans)
        yield non_empty_only(app.translations.get(lang, {}))

CHOICES = {
    'dump-known': DumpKnownAppStrings(commcare_translations.load_translations),
    'simple': SimpleAppStrings(None),
    'select-known': SelectKnownAppStrings(
        functools.partial(commcare_translations.load_translations, version=2)
    ),
}
<|endoftext|>


<commit_before>"""gludb.simple

Provides the simplest possible interface to our functionality.

We provide a simple annotation to create classes with fields (with
optional default values), parameterized constructors, persistence,
and data operations. You are free to derive from any object you wish.
See gludb.data if you need custom or more advanced functionality.

    @DBObject
    class Demo(object):
        some_field = Field()
        my_number = Field(default=42)

    d = Demo(some_field='foo', my_number=3.14)
    print d.to_data()  # Prints a JSON representation
    d1 = Demo.from_data(d.to_data())  # Clone using persistence functions
    d.save()  # Save to database
    for obj in Demo.find_all():  # Print json rep of all objects in DB
        print obj.to_data()
"""

# TODO: key support
# TODO: Indexing support
# TODO: nested objects, lists, lists of nested objects

import json  # TODO: look for other json impls first

from .data import Storable

class Field(object):
    """Support for class-level field declaration.
    """
    def __init__(self, default=''):
        self.name = None
        self.default = default

def _auto_init(self, *args, **kwrds):
    """Our decorator will add this as __init__ to target classes
    """
    for fld in getattr(self, '__fields__', []):
        val = kwrds.get(fld.name, fld.default)
        if not (val is None):
            setattr(self, fld.name, val)

    if callable(getattr(self, 'setup', None)):
        self.setup(*args, **kwrds)

def _get_id(self):
    return self.id

def _set_id(self, new_id):
    self.id = new_id

def _to_data(self):
    data = dict([
        (fld.name, getattr(self, fld.name, fld.default))
        for fld in self.__fields__
    ])
    return json.dumps(data)

def _from_data(cls, data):
    data_dict = json.loads(data)
    return cls(**data_dict)

def DBObject(table_name, versioning):
    """Classes annotated with DBObject gain persistence methods.
    """
    def wrapped(cls):
        field_names = set()
        all_fields = []

        for name in dir(cls):
            fld = getattr(cls, name)
            if fld and isinstance(fld, Field):
                fld.name = name
                all_fields.append(fld)
                field_names.add(name)

        if 'id' not in field_names:
            all_fields.insert(0, Field('id', default=None))

        # Things we count on as part of our processing
        cls.__table_name__ = table_name
        cls.__versioning__ = versioning
        cls.__fields__ = all_fields

        # Give them a ctor for free
        cls.__init__ = _auto_init

        # Duck-type the class for our data methods
        cls.get_id = _get_id
        cls.set_id = _set_id
        cls.to_data = _to_data
        cls.from_data = classmethod(_from_data)

        # Register with our abc since we actually implement all necessary
        # functionality
        Storable.register(cls)

        return cls

    return wrapped
<commit_msg>just use empty str and store instead of None<commit_after>"""gludb.simple

Provides the simplest possible interface to our functionality.

We provide a simple annotation to create classes with fields (with
optional default values), parameterized constructors, persistence,
and data operations. You are free to derive from any object you wish.
See gludb.data if you need custom or more advanced functionality.

    @DBObject
    class Demo(object):
        some_field = Field()
        my_number = Field(default=42)

    d = Demo(some_field='foo', my_number=3.14)
    print d.to_data()  # Prints a JSON representation
    d1 = Demo.from_data(d.to_data())  # Clone using persistence functions
    d.save()  # Save to database
    for obj in Demo.find_all():  # Print json rep of all objects in DB
        print obj.to_data()
"""

# TODO: key support
# TODO: Indexing support
# TODO: nested objects, lists, lists of nested objects

import json  # TODO: look for other json impls first

from .data import Storable

class Field(object):
    """Support for class-level field declaration.
    """
    def __init__(self, default=''):
        self.name = None
        self.default = default

def _auto_init(self, *args, **kwrds):
    """Our decorator will add this as __init__ to target classes
    """
    for fld in getattr(self, '__fields__', []):
        val = kwrds.get(fld.name, fld.default)
        setattr(self, fld.name, val)

    if callable(getattr(self, 'setup', None)):
        self.setup(*args, **kwrds)

def _get_id(self):
    return self.id

def _set_id(self, new_id):
    self.id = new_id

def _to_data(self):
    data = dict([
        (fld.name, getattr(self, fld.name, fld.default))
        for fld in self.__fields__
    ])
    return json.dumps(data)

def _from_data(cls, data):
    data_dict = json.loads(data)
    return cls(**data_dict)

def DBObject(table_name, versioning):
    """Classes annotated with DBObject gain persistence methods.
    """
    def wrapped(cls):
        field_names = set()
        all_fields = []

        for name in dir(cls):
            fld = getattr(cls, name)
            if fld and isinstance(fld, Field):
                fld.name = name
                all_fields.append(fld)
                field_names.add(name)

        if 'id' not in field_names:
            all_fields.insert(0, Field('id', default=''))

        # Things we count on as part of our processing
        cls.__table_name__ = table_name
        cls.__versioning__ = versioning
        cls.__fields__ = all_fields

        # Give them a ctor for free
        cls.__init__ = _auto_init

        # Duck-type the class for our data methods
        cls.get_id = _get_id
        cls.set_id = _set_id
        cls.to_data = _to_data
        cls.from_data = classmethod(_from_data)

        # Register with our abc since we actually implement all necessary
        # functionality
        Storable.register(cls)

        return cls

    return wrapped
<|endoftext|>


<commit_before>def extract_data_fits(filename, use_hdu='all'):
    '''
    Extract non-tabular HDUs from a FITS file. If `use_hdu` is 'all', then
    all non-tabular HDUs are extracted, otherwise only the ones specified
    by `use_hdu` are extracted (`use_hdu` should then contain a list of
    integers). If the requested HDUs do not have the same dimensions, an
    Exception is raised.
    '''
    from ..external.astro import fits

    # Read in all HDUs
    hdulist = fits.open(filename, memmap=True)

    # If only a subset are requested, extract those
    if use_hdu != 'all':
        hdulist = [hdulist[hdu] for hdu in use_hdu]

    # Now only keep HDUs that are not tables
    for hdu in hdulist:
        if not isinstance(hdu, fits.PrimaryHDU) and \
                not isinstance(hdu, fits.ImageHDU):
            hdulist.remove(hdu)

    # Check that dimensions of all HDU are the same
    reference_shape = hdulist[0].data.shape
    for hdu in hdulist:
        if hdu.data.shape != reference_shape:
            raise Exception("HDUs are not all the same dimensions")

    # Extract data
    arrays = {}
    for hdu in hdulist:
        arrays[hdu.name] = hdu.data

    return arrays

def extract_hdf5_datasets(handle):
    '''
    Recursive function that returns a dictionary with all the datasets
    found in an HDF5 file or group. `handle` should be an instance of
    h5py.highlevel.File or h5py.highlevel.Group.
    '''

    import h5py

    datasets = {}
    for group in handle:
        if isinstance(handle[group], h5py.highlevel.Group):
            sub_datasets = extract_hdf5_datasets(handle[group])
            for key in sub_datasets:
                datasets[key] = sub_datasets[key]
        elif isinstance(handle[group], h5py.highlevel.Dataset):
            datasets[handle[group].name] = handle[group]
    return datasets

def extract_data_hdf5(filename, use_datasets='all'):
    '''
    Extract non-tabular datasets from an HDF5 file. If `use_datasets` is
    'all', then all non-tabular datasets are extracted, otherwise only the
    ones specified by `use_datasets` are extracted (`use_datasets` should
    then contain a list of paths). If the requested datasets do not have
    the same dimensions, an Exception is raised.
    '''

    import h5py

    # Open file
    file_handle = h5py.File(filename, 'r')

    # Define function to read

    # Read in all datasets
    datasets = extract_hdf5_datasets(file_handle)

    # Only keep non-tabular datasets
    remove = []
    for key in datasets:
        if datasets[key].dtype.fields is not None:
            remove.append(key)
    for key in remove:
        datasets.pop(key)

    # Check that dimensions of all datasets are the same
    reference_shape = datasets[datasets.keys()[0]].value.shape
    for key in datasets:
        if datasets[key].value.shape != reference_shape:
            raise Exception("Datasets are not all the same dimensions")

    # Extract data
    arrays = {}
    for key in datasets:
        arrays[key] = datasets[key].value

    # Close HDF5 file
    file_handle.close()

    return arrays
<commit_msg>Add new "ignore_blank" keyword to fits.open<commit_after>def extract_data_fits(filename, use_hdu='all'):
    '''
    Extract non-tabular HDUs from a FITS file. If `use_hdu` is 'all', then
    all non-tabular HDUs are extracted, otherwise only the ones specified
    by `use_hdu` are extracted (`use_hdu` should then contain a list of
    integers). If the requested HDUs do not have the same dimensions, an
    Exception is raised.
    '''
    from ..external.astro import fits

    # Read in all HDUs
    hdulist = fits.open(filename, memmap=True, ignore_blank=True)

    # If only a subset are requested, extract those
    if use_hdu != 'all':
        hdulist = [hdulist[hdu] for hdu in use_hdu]

    # Now only keep HDUs that are not tables
    for hdu in hdulist:
        if not isinstance(hdu, fits.PrimaryHDU) and \
                not isinstance(hdu, fits.ImageHDU):
            hdulist.remove(hdu)

    # Check that dimensions of all HDU are the same
    reference_shape = hdulist[0].data.shape
    for hdu in hdulist:
        if hdu.data.shape != reference_shape:
            raise Exception("HDUs are not all the same dimensions")

    # Extract data
    arrays = {}
    for hdu in hdulist:
        arrays[hdu.name] = hdu.data

    return arrays

def extract_hdf5_datasets(handle):
    '''
    Recursive function that returns a dictionary with all the datasets
    found in an HDF5 file or group. `handle` should be an instance of
    h5py.highlevel.File or h5py.highlevel.Group.
    '''

    import h5py

    datasets = {}
    for group in handle:
        if isinstance(handle[group], h5py.highlevel.Group):
            sub_datasets = extract_hdf5_datasets(handle[group])
            for key in sub_datasets:
                datasets[key] = sub_datasets[key]
        elif isinstance(handle[group], h5py.highlevel.Dataset):
            datasets[handle[group].name] = handle[group]
    return datasets

def extract_data_hdf5(filename, use_datasets='all'):
    '''
    Extract non-tabular datasets from an HDF5 file. If `use_datasets` is
    'all', then all non-tabular datasets are extracted, otherwise only the
    ones specified by `use_datasets` are extracted (`use_datasets` should
    then contain a list of paths). If the requested datasets do not have
    the same dimensions, an Exception is raised.
    '''

    import h5py

    # Open file
    file_handle = h5py.File(filename, 'r')

    # Define function to read

    # Read in all datasets
    datasets = extract_hdf5_datasets(file_handle)

    # Only keep non-tabular datasets
    remove = []
    for key in datasets:
        if datasets[key].dtype.fields is not None:
            remove.append(key)
    for key in remove:
        datasets.pop(key)

    # Check that dimensions of all datasets are the same
    reference_shape = datasets[datasets.keys()[0]].value.shape
    for key in datasets:
        if datasets[key].value.shape != reference_shape:
            raise Exception("Datasets are not all the same dimensions")

    # Extract data
    arrays = {}
    for key in datasets:
        arrays[key] = datasets[key].value

    # Close HDF5 file
    file_handle.close()

    return arrays
<|endoftext|>


<commit_before>from __future__ import absolute_import
from __future__ import unicode_literals
from corehq.apps.domain.models import Domain
from corehq.apps.sms.api import send_sms, incoming
from corehq.apps.sms.models import SMS, QueuedSMS
from corehq.apps.sms.tasks import process_sms, MAX_TRIAL_SMS, passes_trial_check
from corehq.apps.sms.tests.util import (
    BaseSMSTest,
    setup_default_sms_test_backend,
    delete_domain_phone_numbers
)
from corehq.apps.smsbillables.models import SmsBillable
from corehq.apps.users.models import CommCareUser
from datetime import datetime, timedelta
from dimagi.utils.couch.cache.cache_core import get_redis_client
from django.conf import settings
from django.test.utils import override_settings
from mock import Mock, patch
from six.moves import range

from corehq.toggles import NAMESPACE_DOMAIN

def patch_datetime_api(timestamp):
    return patch('corehq.apps.sms.api.get_utcnow', new=Mock(return_value=timestamp))

def patch_datetime_tasks(timestamp):
    return patch('corehq.apps.sms.tasks.get_utcnow', new=Mock(return_value=timestamp))

def patch_successful_send():
    return patch('corehq.messaging.smsbackends.test.models.SQLTestSMSBackend.send')

def patch_failed_send():
    return patch(
        'corehq.messaging.smsbackends.test.models.SQLTestSMSBackend.send',
        new=Mock(side_effect=Exception)
    )

def patch_error_send():
    def set_error(msg, *args, **kwargs):
        msg.set_system_error('Error')
    return patch(
        'corehq.messaging.smsbackends.test.models.SQLTestSMSBackend.send',
        new=Mock(side_effect=set_error)
    )

@patch('corehq.apps.sms.management.commands.run_sms_queue.SMSEnqueuingOperation.enqueue', autospec=True)
@patch('corehq.apps.sms.tasks.process_sms.delay', autospec=True)
@override_settings(SMS_QUEUE_ENABLED=True)
class QueueingTestCase(BaseSMSTest):

    def setUp(self):
        super(QueueingTestCase, self).setUp()
        self.domain = 'test-sms-queueing'
        self.domain_obj = Domain(name=self.domain)
        self.domain_obj.save()
        self.create_account_and_subscription(self.domain)
        self.domain_obj = Domain.get(self.domain_obj._id)
        self.backend, self.backend_mapping = setup_default_sms_test_backend()
        self.contact = CommCareUser.create(self.domain, 'user1', 'abc', phone_number='999123')
        entry = self.contact.get_or_create_phone_entry('999123')
        entry.set_two_way()
        entry.set_verified()
        entry.save()

        SmsBillable.objects.filter(domain=self.domain).delete()
        QueuedSMS.objects.all().delete()
        SMS.objects.filter(domain=self.domain).delete()

    def tearDown(self):
        self.contact.delete()
        self.backend.delete()
        self.backend_mapping.delete()

        delete_domain_phone_numbers(self.domain)
        SmsBillable.objects.filter(domain=self.domain).delete()
        QueuedSMS.objects.all().delete()
        SMS.objects.filter(domain=self.domain).delete()
        self.domain_obj.delete()

        super(QueueingTestCase, self).tearDown()

    @property
    def queued_sms_count(self):
        return QueuedSMS.objects.count()

    @property
    def reporting_sms_count(self):
        return SMS.objects.filter(domain=self.domain).count()

    def get_queued_sms(self):
        self.assertEqual(self.queued_sms_count, 1)
        return QueuedSMS.objects.all()[0]

    def get_reporting_sms(self):
        self.assertEqual(self.reporting_sms_count, 1)
        return SMS.objects.filter(domain=self.domain)[0]

    def assertBillableExists(self, msg_id, count=1):
        self.assertEqual(SmsBillable.objects.filter(log_id=msg_id).count(), count)

    def assertBillableDoesNotExist(self, msg_id):
        self.assertEqual(SmsBillable.objects.filter(log_id=msg_id).count(), 0)

    def test_outgoing(self, process_sms_delay_mock, enqueue_directly_mock):
        send_sms(self.domain, None, '+999123', 'test outgoing')

        self.assertEqual(enqueue_directly_mock.call_count, 1)
        self.assertEqual(self.queued_sms_count, 1)
        self.assertEqual(self.reporting_sms_count, 0)

        queued_sms = self.get_queued_sms()
        self.assertEqual(queued_sms.domain, self.domain)
        self.assertEqual(queued_sms.phone_number, '+999123')
        self.assertEqual(queued_sms.text, 'test outgoing')
        self.assertEqual(queued_sms.processed, False)
        self.assertEqual(queued_sms.error, False)
        couch_id = queued_sms.couch_id
        self.assertIsNotNone(couch_id)

        with patch_successful_send() as send_mock:
            process_sms(queued_sms.pk)

        self.assertEqual(send_mock.call_count, 1)
        self.assertEqual(self.queued_sms_count, 0)
        self.assertEqual(self.reporting_sms_count, 1)

        reporting_sms = self.get_reporting_sms()
        self.assertEqual(reporting_sms.domain, self.domain)
        self.assertEqual(reporting_sms.phone_number, '+999123')
        self.assertEqual(reporting_sms.text, 'test outgoing')
        self.assertEqual(reporting_sms.processed, True)
        self.assertEqual(reporting_sms.error, False)
        self.assertEqual(reporting_sms.couch_id, couch_id)
        self.assertEqual(reporting_sms.backend_api, self.backend.get_api_id())
        self.assertEqual(reporting_sms.backend_id, self.backend.couch_id)

        self.assertEqual(process_sms_delay_mock.call_count, 0)
        self.assertBillableExists(couch_id)

    def test_outgoing_failure(self, process_sms_delay_mock, enqueue_directly_mock):
        timestamp = datetime(2016, 1, 1, 12, 0)

        with patch_datetime_api(timestamp):
            send_sms(self.domain, None, '+999123', 'test outgoing')

        self.assertEqual(enqueue_directly_mock.call_count, 1)
        self.assertEqual(self.queued_sms_count, 1)
        self.assertEqual(self.reporting_sms_count, 0)

        for i in range(settings.SMS_QUEUE_MAX_PROCESSING_ATTEMPTS):
            queued_sms = self.get_queued_sms()
            self.assertEqual(queued_sms.domain, self.domain)
            self.assertEqual(queued_sms.phone_number, '+999123')
            self.assertEqual(queued_sms.text, 'test outgoing')
            self.assertEqual(queued_sms.datetime_to_process, timestamp)
            self.assertEqual(queued_sms.processed, False)
            self.assertEqual(queued_sms.error, False)
            self.assertEqual(queued_sms.num_processing_attempts, i)

            with patch_failed_send() as send_mock, patch_datetime_tasks(timestamp + timedelta(seconds=1)):
                process_sms(queued_sms.pk)

            self.assertEqual(process_sms_delay_mock.call_count, 0)
            self.assertBillableDoesNotExist(queued_sms.couch_id)

            self.assertEqual(send_mock.call_count, 1)
            if i < (settings.SMS_QUEUE_MAX_PROCESSING_ATTEMPTS - 1):
                self.assertEqual(self.queued_sms_count, 1)
                self.assertEqual(self.reporting_sms_count, 0)
                timestamp += timedelta(minutes=settings.SMS_QUEUE_REPROCESS_INTERVAL)
            else:
                self.assertEqual(self.queued_sms_count, 0)
                self.assertEqual(self.reporting_sms_count, 1)

        reporting_sms = self.get_reporting_sms()
        self.assertEqual(reporting_sms.domain, self.domain)
        self.assertEqual(reporting_sms.phone_number, '+999123')
        self.assertEqual(reporting_sms.text, 'test outgoing')
        self.assertEqual(reporting_sms.processed, False)
        self.assertEqual(reporting_sms.error, True)
        self.assertEqual(reporting_sms.num_processing_attempts, settings.SMS_QUEUE_MAX_PROCESSING_ATTEMPTS)
        self.assertBillableDoesNotExist(reporting_sms.couch_id)

    def test_outgoing_failure_recovery(self, process_sms_delay_mock, enqueue_directly_mock):
        timestamp = datetime(2016, 1, 1, 12, 0)

        with patch_datetime_api(timestamp):
            send_sms(self.domain, None, '+999123', 'test outgoing')

        self.assertEqual(enqueue_directly_mock.call_count, 1)
        self.assertEqual(self.queued_sms_count, 1)
        self.assertEqual(self.reporting_sms_count, 0)

        queued_sms = self.get_queued_sms()
        self.assertEqual(queued_sms.datetime_to_process, timestamp)
        self.assertEqual(queued_sms.processed, False)
        self.assertEqual(queued_sms.error, False)
        self.assertEqual(queued_sms.num_processing_attempts, 0)

        with patch_failed_send() as send_mock, patch_datetime_tasks(timestamp + timedelta(seconds=1)):
            process_sms(queued_sms.pk)

        self.assertEqual(process_sms_delay_mock.call_count, 0)
        self.assertBillableDoesNotExist(queued_sms.couch_id)
        self.assertEqual(send_mock.call_count, 1)
        self.assertEqual(self.queued_sms_count, 1)
        self.assertEqual(self.reporting_sms_count, 0)
        timestamp += timedelta(minutes=settings.SMS_QUEUE_REPROCESS_INTERVAL)

        queued_sms = self.get_queued_sms()
        self.assertEqual(queued_sms.datetime_to_process, timestamp)
        self.assertEqual(queued_sms.processed, False)
        self.assertEqual(queued_sms.error, False)
        self.assertEqual(queued_sms.num_processing_attempts, 1)

        with patch_successful_send() as send_mock, patch_datetime_tasks(timestamp + timedelta(seconds=1)):
            process_sms(queued_sms.pk)

        self.assertEqual(send_mock.call_count, 1)
        self.assertEqual(self.queued_sms_count, 0)
        self.assertEqual(self.reporting_sms_count, 1)

        reporting_sms = self.get_reporting_sms()
        self.assertEqual(reporting_sms.processed, True)
        self.assertEqual(reporting_sms.error, False)
        self.assertEqual(reporting_sms.num_processing_attempts, 2)

        self.assertEqual(process_sms_delay_mock.call_count, 0)
        self.assertBillableExists(reporting_sms.couch_id)

    def test_outgoing_with_error(self, process_sms_delay_mock, enqueue_directly_mock):
        send_sms(self.domain, None, '+999123', 'test outgoing')

        self.assertEqual(enqueue_directly_mock.call_count, 1)
        self.assertEqual(self.queued_sms_count, 1)
        self.assertEqual(self.reporting_sms_count, 0)

        queued_sms = self.get_queued_sms()
        self.assertEqual(queued_sms.processed, False)
        self.assertEqual(queued_sms.error, False)
        couch_id = queued_sms.couch_id
        self.assertIsNotNone(couch_id)

        with patch_error_send() as send_mock:
            process_sms(queued_sms.pk)

        self.assertEqual(send_mock.call_count, 1)
        self.assertEqual(self.queued_sms_count, 0)
        self.assertEqual(self.reporting_sms_count, 1)

        reporting_sms = self.get_reporting_sms()
        self.assertEqual(reporting_sms.processed, False)
        self.assertEqual(reporting_sms.error, True)
        self.assertEqual(reporting_sms.couch_id, couch_id)
        self.assertEqual(reporting_sms.backend_api, self.backend.get_api_id())
        self.assertEqual(reporting_sms.backend_id, self.backend.couch_id)

        self.assertEqual(process_sms_delay_mock.call_count, 0)
        self.assertBillableDoesNotExist(couch_id)

    @patch.object(QueuedSMS, 'set_system_error')
    @patch('corehq.apps.sms.tasks.domain_is_on_trial')
    def test_passes_trial_check(self, domain_is_on_trial_patch, set_system_error_patch, delay_patch,
            enqueue_patch):

        client = get_redis_client()
        sms = QueuedSMS(domain='trial-project')
        domain_is_on_trial_patch.return_value = True
        key = 'sms-sent-on-trial-for-trial-project'

        # Test when key doesn't exist yet
        self.assertIsNone(client.get(key))
        self.assertTrue(passes_trial_check(sms))
        self.assertEqual(client.get(key), 1)
        self.assertGreater(client.ttl(key), 89 * 24 * 60 * 60)

        # Test with existing key
        self.assertTrue(passes_trial_check(sms))
        self.assertEqual(client.get(key), 2)
        self.assertGreater(client.ttl(key), 89 * 24 * 60 * 60)

        # Test when limit is exceeded
        client.set(key, MAX_TRIAL_SMS)
        self.assertFalse(passes_trial_check(sms))
        set_system_error_patch.assert_called_once()

        # Test when not on trial
        domain_is_on_trial_patch.return_value = False
        self.assertTrue(passes_trial_check(sms))

    def test_incoming(self, process_sms_delay_mock, enqueue_directly_mock):
        incoming('999123', 'inbound test', self.backend.get_api_id())

        self.assertEqual(enqueue_directly_mock.call_count, 1)
        self.assertEqual(self.queued_sms_count, 1)
        self.assertEqual(self.reporting_sms_count, 0)

        queued_sms = self.get_queued_sms()
        self.assertIsNone(queued_sms.domain)
        self.assertIsNone(queued_sms.couch_recipient_doc_type)
        self.assertIsNone(queued_sms.couch_recipient)
        self.assertEqual(queued_sms.phone_number, '+999123')
        self.assertEqual(queued_sms.text, 'inbound test')
        self.assertEqual(queued_sms.processed, False)
        self.assertEqual(queued_sms.error, False)
        self.assertEqual(queued_sms.backend_api, self.backend.get_api_id())
        couch_id = queued_sms.couch_id
        self.assertIsNotNone(couch_id)
        self.assertBillableDoesNotExist(couch_id)

        process_sms(queued_sms.pk)
        self.assertEqual(self.queued_sms_count, 0)
        self.assertEqual(self.reporting_sms_count, 1)

        reporting_sms = self.get_reporting_sms()
        self.assertEqual(reporting_sms.domain, self.domain)
        self.assertEqual(reporting_sms.couch_recipient_doc_type, self.contact.doc_type)
        self.assertEqual(reporting_sms.couch_recipient, self.contact.get_id)
        self.assertEqual(reporting_sms.phone_number, '+999123')
        self.assertEqual(reporting_sms.text, 'inbound test')
        self.assertEqual(reporting_sms.processed, True)
        self.assertEqual(reporting_sms.error, False)
        self.assertEqual(reporting_sms.backend_api, self.backend.get_api_id())
        self.assertEqual(reporting_sms.couch_id, couch_id)
        self.assertBillableExists(couch_id)
<commit_msg>Fixed lint error<commit_after>from __future__ import absolute_import
from __future__ import unicode_literals
from corehq.apps.domain.models import Domain
from corehq.apps.sms.api import send_sms, incoming
from corehq.apps.sms.models import SMS, QueuedSMS
from corehq.apps.sms.tasks import process_sms, MAX_TRIAL_SMS, passes_trial_check
from corehq.apps.sms.tests.util import (
    BaseSMSTest,
    setup_default_sms_test_backend,
    delete_domain_phone_numbers
)
from corehq.apps.smsbillables.models import SmsBillable
from corehq.apps.users.models import CommCareUser
from datetime import datetime, timedelta
from dimagi.utils.couch.cache.cache_core import get_redis_client
from django.conf import settings
from django.test.utils import override_settings
from mock import Mock, patch
from six.moves import range

def patch_datetime_api(timestamp):
    return patch('corehq.apps.sms.api.get_utcnow', new=Mock(return_value=timestamp))

def patch_datetime_tasks(timestamp):
    return patch('corehq.apps.sms.tasks.get_utcnow', new=Mock(return_value=timestamp))

def patch_successful_send():
    return patch('corehq.messaging.smsbackends.test.models.SQLTestSMSBackend.send')

def patch_failed_send():
    return patch(
        'corehq.messaging.smsbackends.test.models.SQLTestSMSBackend.send',
        new=Mock(side_effect=Exception)
    )

def patch_error_send():
    def set_error(msg, *args, **kwargs):
        msg.set_system_error('Error')
    return patch(
        'corehq.messaging.smsbackends.test.models.SQLTestSMSBackend.send',
        new=Mock(side_effect=set_error)
    )

@patch('corehq.apps.sms.management.commands.run_sms_queue.SMSEnqueuingOperation.enqueue', autospec=True)
@patch('corehq.apps.sms.tasks.process_sms.delay', autospec=True)
@override_settings(SMS_QUEUE_ENABLED=True)
class QueueingTestCase(BaseSMSTest):

    def setUp(self):
        super(QueueingTestCase, self).setUp()
        self.domain = 'test-sms-queueing'
        self.domain_obj = Domain(name=self.domain)
        self.domain_obj.save()
        self.create_account_and_subscription(self.domain)
        self.domain_obj = Domain.get(self.domain_obj._id)
        self.backend, self.backend_mapping = setup_default_sms_test_backend()
        self.contact = CommCareUser.create(self.domain, 'user1', 'abc', phone_number='999123')
        entry = self.contact.get_or_create_phone_entry('999123')
        entry.set_two_way()
        entry.set_verified()
        entry.save()

        SmsBillable.objects.filter(domain=self.domain).delete()
        QueuedSMS.objects.all().delete()
        SMS.objects.filter(domain=self.domain).delete()

    def tearDown(self):
        self.contact.delete()
        self.backend.delete()
        self.backend_mapping.delete()

        delete_domain_phone_numbers(self.domain)
        SmsBillable.objects.filter(domain=self.domain).delete()
        QueuedSMS.objects.all().delete()
        SMS.objects.filter(domain=self.domain).delete()
        self.domain_obj.delete()

        super(QueueingTestCase, self).tearDown()

    @property
    def queued_sms_count(self):
        return QueuedSMS.objects.count()

    @property
    def reporting_sms_count(self):
        return SMS.objects.filter(domain=self.domain).count()

    def get_queued_sms(self):
        self.assertEqual(self.queued_sms_count, 1)
        return QueuedSMS.objects.all()[0]

    def get_reporting_sms(self):
        self.assertEqual(self.reporting_sms_count, 1)
        return SMS.objects.filter(domain=self.domain)[0]

    def assertBillableExists(self, msg_id, count=1):
        self.assertEqual(SmsBillable.objects.filter(log_id=msg_id).count(), count)

    def assertBillableDoesNotExist(self, msg_id):
        self.assertEqual(SmsBillable.objects.filter(log_id=msg_id).count(), 0)

    def test_outgoing(self, process_sms_delay_mock, enqueue_directly_mock):
        send_sms(self.domain, None, '+999123', 'test outgoing')

        self.assertEqual(enqueue_directly_mock.call_count, 1)
        self.assertEqual(self.queued_sms_count, 1)
        self.assertEqual(self.reporting_sms_count, 0)

        queued_sms = self.get_queued_sms()
        self.assertEqual(queued_sms.domain, self.domain)
        self.assertEqual(queued_sms.phone_number, '+999123')
        self.assertEqual(queued_sms.text, 'test outgoing')
        self.assertEqual(queued_sms.processed, False)
        self.assertEqual(queued_sms.error, False)
        couch_id = queued_sms.couch_id
        self.assertIsNotNone(couch_id)

        with patch_successful_send() as send_mock:
            process_sms(queued_sms.pk)

        self.assertEqual(send_mock.call_count, 1)
        self.assertEqual(self.queued_sms_count, 0)
        self.assertEqual(self.reporting_sms_count, 1)

        reporting_sms = self.get_reporting_sms()
        self.assertEqual(reporting_sms.domain, self.domain)
        self.assertEqual(reporting_sms.phone_number, '+999123')
        self.assertEqual(reporting_sms.text, 'test outgoing')
        self.assertEqual(reporting_sms.processed, True)
        self.assertEqual(reporting_sms.error, False)
        self.assertEqual(reporting_sms.couch_id, couch_id)
        self.assertEqual(reporting_sms.backend_api, self.backend.get_api_id())
        self.assertEqual(reporting_sms.backend_id, self.backend.couch_id)

        self.assertEqual(process_sms_delay_mock.call_count, 0)
        self.assertBillableExists(couch_id)

    def test_outgoing_failure(self, process_sms_delay_mock, enqueue_directly_mock):
        timestamp = datetime(2016, 1, 1, 12, 0)

        with patch_datetime_api(timestamp):
            send_sms(self.domain, None, '+999123', 'test outgoing')

        self.assertEqual(enqueue_directly_mock.call_count, 1)
        self.assertEqual(self.queued_sms_count, 1)
        self.assertEqual(self.reporting_sms_count, 0)

        for i in range(settings.SMS_QUEUE_MAX_PROCESSING_ATTEMPTS):
            queued_sms = self.get_queued_sms()
            self.assertEqual(queued_sms.domain, self.domain)
            self.assertEqual(queued_sms.phone_number, '+999123')
            self.assertEqual(queued_sms.text, 'test outgoing')
            self.assertEqual(queued_sms.datetime_to_process, timestamp)
            self.assertEqual(queued_sms.processed, False)
            self.assertEqual(queued_sms.error, False)
            self.assertEqual(queued_sms.num_processing_attempts, i)

            with patch_failed_send() as send_mock, patch_datetime_tasks(timestamp + timedelta(seconds=1)):
                process_sms(queued_sms.pk)

            self.assertEqual(process_sms_delay_mock.call_count, 0)
            self.assertBillableDoesNotExist(queued_sms.couch_id)

            self.assertEqual(send_mock.call_count, 1)
            if i < (settings.SMS_QUEUE_MAX_PROCESSING_ATTEMPTS - 1):
                self.assertEqual(self.queued_sms_count, 1)
                self.assertEqual(self.reporting_sms_count, 0)
                timestamp += timedelta(minutes=settings.SMS_QUEUE_REPROCESS_INTERVAL)
            else:
                self.assertEqual(self.queued_sms_count, 0)
                self.assertEqual(self.reporting_sms_count, 1)

        reporting_sms = self.get_reporting_sms()
        self.assertEqual(reporting_sms.domain, self.domain)
        self.assertEqual(reporting_sms.phone_number, '+999123')
        self.assertEqual(reporting_sms.text, 'test outgoing')
        self.assertEqual(reporting_sms.processed, False)
        self.assertEqual(reporting_sms.error, True)
        self.assertEqual(reporting_sms.num_processing_attempts, settings.SMS_QUEUE_MAX_PROCESSING_ATTEMPTS)
        self.assertBillableDoesNotExist(reporting_sms.couch_id)

    def test_outgoing_failure_recovery(self, process_sms_delay_mock, enqueue_directly_mock):
        timestamp = datetime(2016, 1, 1, 12, 0)

        with patch_datetime_api(timestamp):
            send_sms(self.domain, None, '+999123', 'test outgoing')

        self.assertEqual(enqueue_directly_mock.call_count, 1)
        self.assertEqual(self.queued_sms_count, 1)
        self.assertEqual(self.reporting_sms_count, 0)

        queued_sms = self.get_queued_sms()
        self.assertEqual(queued_sms.datetime_to_process, timestamp)
        self.assertEqual(queued_sms.processed, False)
        self.assertEqual(queued_sms.error, False)
        self.assertEqual(queued_sms.num_processing_attempts, 0)

        with patch_failed_send() as send_mock, patch_datetime_tasks(timestamp + timedelta(seconds=1)):
            process_sms(queued_sms.pk)

        self.assertEqual(process_sms_delay_mock.call_count, 0)
        self.assertBillableDoesNotExist(queued_sms.couch_id)
        self.assertEqual(send_mock.call_count, 1)
        self.assertEqual(self.queued_sms_count, 1)
        self.assertEqual(self.reporting_sms_count, 0)
        timestamp += timedelta(minutes=settings.SMS_QUEUE_REPROCESS_INTERVAL)

        queued_sms = self.get_queued_sms()
        self.assertEqual(queued_sms.datetime_to_process, timestamp)
        self.assertEqual(queued_sms.processed, False)
        self.assertEqual(queued_sms.error, False)
        self.assertEqual(queued_sms.num_processing_attempts, 1)

        with patch_successful_send() as send_mock, patch_datetime_tasks(timestamp + timedelta(seconds=1)):
            process_sms(queued_sms.pk)

        self.assertEqual(send_mock.call_count, 1)
        self.assertEqual(self.queued_sms_count, 0)
        self.assertEqual(self.reporting_sms_count, 1)

        reporting_sms = self.get_reporting_sms()
        self.assertEqual(reporting_sms.processed, True)
        self.assertEqual(reporting_sms.error, False)
        self.assertEqual(reporting_sms.num_processing_attempts, 2)

        self.assertEqual(process_sms_delay_mock.call_count, 0)
        self.assertBillableExists(reporting_sms.couch_id)

    def test_outgoing_with_error(self, process_sms_delay_mock, enqueue_directly_mock):
        send_sms(self.domain, None, '+999123', 'test outgoing')

        self.assertEqual(enqueue_directly_mock.call_count, 1)
        self.assertEqual(self.queued_sms_count, 1)
        self.assertEqual(self.reporting_sms_count, 0)

        queued_sms = self.get_queued_sms()
        self.assertEqual(queued_sms.processed, False)
        self.assertEqual(queued_sms.error, False)
        couch_id = queued_sms.couch_id
        self.assertIsNotNone(couch_id)

        with patch_error_send() as send_mock:
            process_sms(queued_sms.pk)

        self.assertEqual(send_mock.call_count, 1)
        self.assertEqual(self.queued_sms_count, 0)
        self.assertEqual(self.reporting_sms_count, 1)

        reporting_sms = self.get_reporting_sms()
        self.assertEqual(reporting_sms.processed, False)
        self.assertEqual(reporting_sms.error, True)
        self.assertEqual(reporting_sms.couch_id, couch_id)
        self.assertEqual(reporting_sms.backend_api, self.backend.get_api_id())
        self.assertEqual(reporting_sms.backend_id, self.backend.couch_id)

        self.assertEqual(process_sms_delay_mock.call_count, 0)
        self.assertBillableDoesNotExist(couch_id)

    @patch.object(QueuedSMS, 'set_system_error')
    @patch('corehq.apps.sms.tasks.domain_is_on_trial')
    def test_passes_trial_check(self, domain_is_on_trial_patch, set_system_error_patch, delay_patch,
            enqueue_patch):

        client = get_redis_client()
        sms = QueuedSMS(domain='trial-project')
        domain_is_on_trial_patch.return_value = True
        key = 'sms-sent-on-trial-for-trial-project'

        # Test when key doesn't exist yet
        self.assertIsNone(client.get(key))
        self.assertTrue(passes_trial_check(sms))
        self.assertEqual(client.get(key), 1)
        self.assertGreater(client.ttl(key), 89 * 24 * 60 * 60)

        # Test with existing key
        self.assertTrue(passes_trial_check(sms))
        self.assertEqual(client.get(key), 2)
        self.assertGreater(client.ttl(key), 89 * 24 * 60 * 60)

        # Test when limit is exceeded
        client.set(key, MAX_TRIAL_SMS)
        self.assertFalse(passes_trial_check(sms))
        set_system_error_patch.assert_called_once()

        # Test when not on trial
        domain_is_on_trial_patch.return_value = False
        self.assertTrue(passes_trial_check(sms))

    def test_incoming(self, process_sms_delay_mock, enqueue_directly_mock):
        incoming('999123', 'inbound test', self.backend.get_api_id())

        self.assertEqual(enqueue_directly_mock.call_count, 1)
        self.assertEqual(self.queued_sms_count, 1)
        self.assertEqual(self.reporting_sms_count, 0)

        queued_sms = self.get_queued_sms()
        self.assertIsNone(queued_sms.domain)
        self.assertIsNone(queued_sms.couch_recipient_doc_type)
        self.assertIsNone(queued_sms.couch_recipient)
        self.assertEqual(queued_sms.phone_number, '+999123')
        self.assertEqual(queued_sms.text, 'inbound test')
        self.assertEqual(queued_sms.processed, False)
        self.assertEqual(queued_sms.error, False)
        self.assertEqual(queued_sms.backend_api, self.backend.get_api_id())
        couch_id = queued_sms.couch_id
        self.assertIsNotNone(couch_id)
        self.assertBillableDoesNotExist(couch_id)

        process_sms(queued_sms.pk)
        self.assertEqual(self.queued_sms_count, 0)
        self.assertEqual(self.reporting_sms_count, 1)

        reporting_sms = self.get_reporting_sms()
        self.assertEqual(reporting_sms.domain, self.domain)
        self.assertEqual(reporting_sms.couch_recipient_doc_type, self.contact.doc_type)
        self.assertEqual(reporting_sms.couch_recipient, self.contact.get_id)
        self.assertEqual(reporting_sms.phone_number, '+999123')
        self.assertEqual(reporting_sms.text, 'inbound test')
        self.assertEqual(reporting_sms.processed, True)
        self.assertEqual(reporting_sms.error, False)
        self.assertEqual(reporting_sms.backend_api, self.backend.get_api_id())
        self.assertEqual(reporting_sms.couch_id, couch_id)
        self.assertBillableExists(couch_id)
<|endoftext|>


<commit_before>from django.contrib import auth
from django.test import TestCase, RequestFactory

from corehq.apps.domain.models import Domain
from corehq.apps.sso.models import IdentityProvider, AuthenticatedEmailDomain
from corehq.apps.sso.tests import generator
from corehq.apps.users.models import WebUser

class TestSsoBackend(TestCase):
    """
    This ensures that authentication through the SsoBackend only succeeds
    when a samlSessionIndex is present and the Identity Provider, its
    associated email domains, and the username all meet required criteria.
    """

    @classmethod
    def setUpClass(cls):
        super().setUpClass()
        cls.account = generator.get_billing_account_for_idp()
        cls.domain = Domain.get_or_create_with_name("vaultwax-001", is_active=True)

        # this will be the user that's "logging in" with SAML2 via the SsoBackend
        cls.user = WebUser.create(
            cls.domain.name, 'b@vaultwax.com', 'testpwd', None, None
        )
        cls.idp = generator.create_idp('vaultwax', cls.account)
        cls.idp_two = generator.create_idp('vwx-link', cls.account)

    @classmethod
    def tearDownClass(cls):
        IdentityProvider.objects.all().delete()
        cls.user.delete(None)
        cls.domain.delete()
        cls.account.delete()
        super().tearDownClass()

    def setUp(self):
        super().setUp()
        # Identity Provider is_active is always True by default for each test
        self.idp.is_active = True
        self.idp.save()

        self.request = RequestFactory().get('/sso/test')
        self.request.session = {
            'samlSessionIndex': '_7c84c96e-8774-4e64-893c-06f91d285100',
        }

    def tearDown(self):
        AuthenticatedEmailDomain.objects.all().delete()
        super().tearDown()

    def test_backend_failure_without_idp_slug(self):
        """
        SsoBackend should fail to move past the first check because an idp_slug
         was not passed to authenticate()
        """
        user = auth.authenticate(
            request=self.request,
            username=self.user.username
        )
        self.assertIsNone(user)
        with self.assertRaises(AttributeError):
            self.request.sso_login_error

    def test_backend_failure_without_saml_session_index(self):
        """
        SsoBackend should fail to move past the first check because a
         samlSessionIndex is not present in request.session.
        """
        del self.request.session['samlSessionIndex']
        user = auth.authenticate(
            request=self.request,
            username=self.user.username,
            idp_slug=self.idp.slug
        )
        self.assertIsNone(user)

    def test_login_error_if_idp_doesnt_exist(self):
        """
        SsoBackend should fail to return a user if the Identity Provider
        slug associated with that user does not exist. It should also populate
        request.sso_login_error with an error message.
        """
        user = auth.authenticate(
            request=self.request,
            username=self.user.username,
            idp_slug='doesnotexist'
        )
        self.assertIsNone(user)
        self.assertEqual(
            self.request.sso_login_error,
            "Identity Provider does not exist."
        )

    def test_login_error_if_idp_not_active(self):
        """
        SsoBackend should fail to return a user if the Identity Provider
        associated with that user is not active. It should also populate
        request.sso_login_error with an error message.
        """
        self.idp.is_active = False
        self.idp.save()
        user = auth.authenticate(
            request=self.request,
            username=self.user.username,
            idp_slug=self.idp.slug
        )
        self.assertIsNone(user)
        self.assertEqual(
            self.request.sso_login_error,
            "This Identity Provider is not active."
        )

    def test_login_error_if_bad_username(self):
        """
        SsoBackend should fail to return a user if the username passed to it
        is not a valid email address (missing email domain / no `@`). It should
        also populate request.sso_login_error with an error message.
        """
        user = auth.authenticate(
            request=self.request,
            username='badusername',
            idp_slug=self.idp.slug
        )
        self.assertIsNone(user)
        self.assertEqual(
            self.request.sso_login_error,
            "Username badusername is not valid."
        )

    def test_login_error_if_email_domain_does_not_exist(self):
        """
        SsoBackend should fail to return a user if the username passed to it
        matches an email domain that is not mapped to any Identity Provider.
        It should also populate request.sso_login_error with an error message.
        """
        user = auth.authenticate(
            request=self.request,
            username='b@vwx.link',
            idp_slug=self.idp.slug
        )
        self.assertIsNone(user)
        self.assertEqual(
            self.request.sso_login_error,
            "The Email Domain vwx.link is not allowed to authenticate with "
            "this Identity Provider."
        )

    def test_login_error_if_email_domain_is_not_authorized(self):
        """
        SsoBackend should fail to return a user if the username passed to it
        matches an email domain that exists, but is not authorized to
        authenticate with the Identity Provider in the request. It should also
        populate request.sso_login_error with an error message.
        """
        AuthenticatedEmailDomain.objects.create(
            email_domain='vwx.link',
            identity_provider=self.idp_two,  # note that this is not self.idp
        )
        user = auth.authenticate(
            request=self.request,
            username='b@vwx.link',
            idp_slug=self.idp.slug  # note that this is self.idp, not idp_two
        )
        self.assertIsNone(user)
        self.assertEqual(
            self.request.sso_login_error,
            "The Email Domain vwx.link is not allowed to authenticate with "
            "this Identity Provider."
        )

    def test_login_error_if_user_does_not_exist(self):
        """
        SsoBackend should fail to return a user if the username passed to does
        not exist. It should also populate request.sso_login_error with an
        error message.

        todo this test will change with additional user creation workflows
         that will be introduced later
        """
        AuthenticatedEmailDomain.objects.create(
            email_domain='vaultwax.com',
            identity_provider=self.idp,
        )
        user = auth.authenticate(
            request=self.request,
            username='testnoexist@vaultwax.com',
            idp_slug=self.idp.slug
        )
        self.assertIsNone(user)
        self.assertEqual(
            self.request.sso_login_error,
            "User testnoexist@vaultwax.com does not exist."
        )

    def test_successful_login(self):
        """
        This test demonstrates the requirements necessary for a SsoBackend to
        successfully return a user and report no login error.
        """
        AuthenticatedEmailDomain.objects.create(
            email_domain='vaultwax.com',
            identity_provider=self.idp,
        )
        user = auth.authenticate(
            request=self.request,
            username=self.user.username,
            idp_slug=self.idp.slug
        )
        self.assertIsNotNone(user)
        self.assertEqual(user.username, self.user.username)
        self.assertIsNone(self.request.sso_login_error)
<commit_msg>clarify comment<commit_after>from django.contrib import auth
from django.test import TestCase, RequestFactory

from corehq.apps.domain.models import Domain
from corehq.apps.sso.models import IdentityProvider, AuthenticatedEmailDomain
from corehq.apps.sso.tests import generator
from corehq.apps.users.models import WebUser

class TestSsoBackend(TestCase):
    """
    This ensures that authentication through the SsoBackend only succeeds
    when a samlSessionIndex is present and the Identity Provider, its
    associated email domains, and the username all meet required criteria.
    """

    @classmethod
    def setUpClass(cls):
        super().setUpClass()
        cls.account = generator.get_billing_account_for_idp()
        cls.domain = Domain.get_or_create_with_name("vaultwax-001", is_active=True)

        # this will be the user that's "logging in" with SAML2 via the SsoBackend
        cls.user = WebUser.create(
            cls.domain.name, 'b@vaultwax.com', 'testpwd', None, None
        )
        cls.idp = generator.create_idp('vaultwax', cls.account)
        cls.idp_two = generator.create_idp('vwx-link', cls.account)

    @classmethod
    def tearDownClass(cls):
        IdentityProvider.objects.all().delete()
        cls.user.delete(None)
        cls.domain.delete()
        cls.account.delete()
        super().tearDownClass()

    def setUp(self):
        super().setUp()
        # some tests set is_active to False, so make sure is_active is always
        #  True by default
        self.idp.is_active = True
        self.idp.save()

        self.request = RequestFactory().get('/sso/test')
        self.request.session = {
            'samlSessionIndex': '_7c84c96e-8774-4e64-893c-06f91d285100',
        }

    def tearDown(self):
        AuthenticatedEmailDomain.objects.all().delete()
        super().tearDown()

    def test_backend_failure_without_idp_slug(self):
        """
        SsoBackend should fail to move past the first check because an idp_slug
         was not passed to authenticate()
        """
        user = auth.authenticate(
            request=self.request,
            username=self.user.username
        )
        self.assertIsNone(user)
        with self.assertRaises(AttributeError):
            self.request.sso_login_error

    def test_backend_failure_without_saml_session_index(self):
        """
        SsoBackend should fail to move past the first check because a
         samlSessionIndex is not present in request.session.
        """
        del self.request.session['samlSessionIndex']
        user = auth.authenticate(
            request=self.request,
            username=self.user.username,
            idp_slug=self.idp.slug
        )
        self.assertIsNone(user)

    def test_login_error_if_idp_doesnt_exist(self):
        """
        SsoBackend should fail to return a user if the Identity Provider
        slug associated with that user does not exist. It should also populate
        request.sso_login_error with an error message.
        """
        user = auth.authenticate(
            request=self.request,
            username=self.user.username,
            idp_slug='doesnotexist'
        )
        self.assertIsNone(user)
        self.assertEqual(
            self.request.sso_login_error,
            "Identity Provider does not exist."
        )

    def test_login_error_if_idp_not_active(self):
        """
        SsoBackend should fail to return a user if the Identity Provider
        associated with that user is not active. It should also populate
        request.sso_login_error with an error message.
        """
        self.idp.is_active = False
        self.idp.save()
        user = auth.authenticate(
            request=self.request,
            username=self.user.username,
            idp_slug=self.idp.slug
        )
        self.assertIsNone(user)
        self.assertEqual(
            self.request.sso_login_error,
            "This Identity Provider is not active."
        )

    def test_login_error_if_bad_username(self):
        """
        SsoBackend should fail to return a user if the username passed to it
        is not a valid email address (missing email domain / no `@`). It should
        also populate request.sso_login_error with an error message.
        """
        user = auth.authenticate(
            request=self.request,
            username='badusername',
            idp_slug=self.idp.slug
        )
        self.assertIsNone(user)
        self.assertEqual(
            self.request.sso_login_error,
            "Username badusername is not valid."
        )

    def test_login_error_if_email_domain_does_not_exist(self):
        """
        SsoBackend should fail to return a user if the username passed to it
        matches an email domain that is not mapped to any Identity Provider.
        It should also populate request.sso_login_error with an error message.
        """
        user = auth.authenticate(
            request=self.request,
            username='b@vwx.link',
            idp_slug=self.idp.slug
        )
        self.assertIsNone(user)
        self.assertEqual(
            self.request.sso_login_error,
            "The Email Domain vwx.link is not allowed to authenticate with "
            "this Identity Provider."
        )

    def test_login_error_if_email_domain_is_not_authorized(self):
        """
        SsoBackend should fail to return a user if the username passed to it
        matches an email domain that exists, but is not authorized to
        authenticate with the Identity Provider in the request. It should also
        populate request.sso_login_error with an error message.
        """
        AuthenticatedEmailDomain.objects.create(
            email_domain='vwx.link',
            identity_provider=self.idp_two,  # note that this is not self.idp
        )
        user = auth.authenticate(
            request=self.request,
            username='b@vwx.link',
            idp_slug=self.idp.slug  # note that this is self.idp, not idp_two
        )
        self.assertIsNone(user)
        self.assertEqual(
            self.request.sso_login_error,
            "The Email Domain vwx.link is not allowed to authenticate with "
            "this Identity Provider."
        )

    def test_login_error_if_user_does_not_exist(self):
        """
        SsoBackend should fail to return a user if the username passed to does
        not exist. It should also populate request.sso_login_error with an
        error message.

        todo this test will change with additional user creation workflows
         that will be introduced later
        """
        AuthenticatedEmailDomain.objects.create(
            email_domain='vaultwax.com',
            identity_provider=self.idp,
        )
        user = auth.authenticate(
            request=self.request,
            username='testnoexist@vaultwax.com',
            idp_slug=self.idp.slug
        )
        self.assertIsNone(user)
        self.assertEqual(
            self.request.sso_login_error,
            "User testnoexist@vaultwax.com does not exist."
        )

    def test_successful_login(self):
        """
        This test demonstrates the requirements necessary for a SsoBackend to
        successfully return a user and report no login error.
        """
        AuthenticatedEmailDomain.objects.create(
            email_domain='vaultwax.com',
            identity_provider=self.idp,
        )
        user = auth.authenticate(
            request=self.request,
            username=self.user.username,
            idp_slug=self.idp.slug
        )
        self.assertIsNotNone(user)
        self.assertEqual(user.username, self.user.username)
        self.assertIsNone(self.request.sso_login_error)
<|endoftext|>


<commit_before>import time, urllib2, base64
from bs4 import BeautifulSoup as bs

config = ConfigParser.ConfigParser()
config.read("utorrent.config")
url = config.get("webui","url") + '/gui/'
authString = base64.b64encode(config.get("webui","username") + ':' + config.get("webui","password"))

class Utorrent:
    def __init__(self):
       self.authRequest = Utorrent.doAuthRequest()
       self.token = self.getToken()
       self.cookie = self.getCookie()

    @staticmethod
    def doAuthRequest():
        url = webUiUrl + 'token.html?t=' + str(int(time.time()))
        request = urllib2.Request(url)
        request.add_header('Authorization', 'Basic ' + authString)
        return urllib2.urlopen(request)

    def getToken(self):
        parsedAuthRequest = bs(self.authRequest)
        return parsedAuthRequest.div.get_text()

    def getCookie(self):
        return self.authRequest.info()['Set-Cookie']

    def download(self, url):
        downloadUrl = webUiUrl + '?token='+ self.token + '&action=add-url&s=' + url
        print downloadUrl
        addTorrentRequest = urllib2.Request(downloadUrl)

        addTorrentRequest.add_header('Cookie', self.cookie)
        addTorrentRequest.add_header('Authorization', 'Basic ' + authString)
        result = ""
        try:
            result = urllib2.urlopen(addTorrentRequest).read()
        except urllib2.HTTPError, err:
            result = err
        return result
<commit_msg>add missing dependency to ConfigParser<commit_after>import time, urllib2, base64, ConfigParser
from bs4 import BeautifulSoup as bs

config = ConfigParser.ConfigParser()
config.read("utorrent.config")
url = config.get("webui","url") + '/gui/'
authString = base64.b64encode(config.get("webui","username") + ':' + config.get("webui","password"))

class Utorrent:
    def __init__(self):
       self.authRequest = Utorrent.doAuthRequest()
       self.token = self.getToken()
       self.cookie = self.getCookie()

    @staticmethod
    def doAuthRequest():
        url = webUiUrl + 'token.html?t=' + str(int(time.time()))
        request = urllib2.Request(url)
        request.add_header('Authorization', 'Basic ' + authString)
        return urllib2.urlopen(request)

    def getToken(self):
        parsedAuthRequest = bs(self.authRequest)
        return parsedAuthRequest.div.get_text()

    def getCookie(self):
        return self.authRequest.info()['Set-Cookie']

    def download(self, url):
        downloadUrl = webUiUrl + '?token='+ self.token + '&action=add-url&s=' + url
        print downloadUrl
        addTorrentRequest = urllib2.Request(downloadUrl)

        addTorrentRequest.add_header('Cookie', self.cookie)
        addTorrentRequest.add_header('Authorization', 'Basic ' + authString)
        result = ""
        try:
            result = urllib2.urlopen(addTorrentRequest).read()
        except urllib2.HTTPError, err:
            result = err
        return result
<|endoftext|>


<commit_before>from django.conf import settings

from corehq.toggles import USE_SQL_BACKEND, NAMESPACE_DOMAIN, NEW_EXPORTS, TF_USES_SQLITE_BACKEND
from dimagi.utils.logging import notify_exception

def should_use_sql_backend(domain_object_or_name):
    from corehq.apps.domain.models import Domain
    if settings.UNIT_TESTING:
        return _should_use_sql_backend_in_tests(domain_object_or_name)

    # TODO: remove toggle once all domains have been migrated
    if isinstance(domain_object_or_name, Domain):
        domain_name = domain_object_or_name.name
        domain_object = domain_object_or_name
    else:
        domain_name = domain_object_or_name
        domain_object = Domain.get_by_name(domain_name)

    if domain_object is None:
        return False
    toggle_enabled = USE_SQL_BACKEND.enabled(domain_name)
    if toggle_enabled:
        try:
            # migrate domains in toggle
            if not domain_object.use_sql_backend:
                domain_object.use_sql_backend = True
                domain_object.save()
                USE_SQL_BACKEND.set(domain_name, enabled=False, namespace=NAMESPACE_DOMAIN)
        except Exception:
            notify_exception(None, "Error migrating SQL BACKEND toggle", {
                'domain': domain_name
            })
        return True
    else:
        return domain_object.use_sql_backend

def _should_use_sql_backend_in_tests(domain_object_or_name):
    """The default return value is False unless the ``TESTS_SHOULD_USE_SQL_BACKEND`` setting
    has been set or a Domain object with the same name exists."""
    assert settings.UNIT_TESTING
    from corehq.apps.domain.models import Domain
    override = getattr(settings, 'TESTS_SHOULD_USE_SQL_BACKEND', None)
    if override is not None:
        return override

    if domain_object_or_name and getattr(settings, 'DB_ENABLED', True):
        domain_object = domain_object_or_name \
            if isinstance(domain_object_or_name, Domain) \
            else Domain.get_by_name(domain_object_or_name)
        return domain_object and domain_object.use_sql_backend
    else:
        return False

def use_new_exports(domain_name):
    return NEW_EXPORTS.enabled(domain_name) or should_use_sql_backend(domain_name)

def use_sqlite_backend(domain_name):
    return TF_USES_SQLITE_BACKEND.enabled(domain_name) or should_use_sql_backend(domain_name)

def is_commcarecase(obj):
    from casexml.apps.case.models import CommCareCase
    from corehq.form_processor.models import CommCareCaseSQL
    return isinstance(obj, (CommCareCase, CommCareCaseSQL))
<commit_msg>add thread local override for sql backend setting<commit_after>from django.conf import settings

from corehq.toggles import USE_SQL_BACKEND, NAMESPACE_DOMAIN, NEW_EXPORTS, TF_USES_SQLITE_BACKEND
from dimagi.utils.logging import notify_exception

import threading

_thread_local = threading.local()

def get_local_domain_sql_backend_override(domain):
    try:
        return _thread_local.use_sql_backend[domain]
    except (AttributeError, KeyError):
        return None

def set_local_domain_sql_backend_override(domain):
    use_sql_backend_dict = getattr(_thread_local, 'use_sql_backend', {})
    use_sql_backend_dict[domain] = True
    _thread_local.use_sql_backend = use_sql_backend_dict

def should_use_sql_backend(domain_object_or_name):
    from corehq.apps.domain.models import Domain
    if settings.UNIT_TESTING:
        return _should_use_sql_backend_in_tests(domain_object_or_name)

    # TODO: remove toggle once all domains have been migrated
    if isinstance(domain_object_or_name, Domain):
        domain_name = domain_object_or_name.name
        domain_object = domain_object_or_name
    else:
        domain_name = domain_object_or_name
        domain_object = Domain.get_by_name(domain_name)

    if domain_object is None:
        return False

    local_override = get_local_domain_sql_backend_override(domain_name)
    if local_override is not None:
        return local_override

    toggle_enabled = USE_SQL_BACKEND.enabled(domain_name)
    if toggle_enabled:
        try:
            # migrate domains in toggle
            if not domain_object.use_sql_backend:
                domain_object.use_sql_backend = True
                domain_object.save()
                USE_SQL_BACKEND.set(domain_name, enabled=False, namespace=NAMESPACE_DOMAIN)
        except Exception:
            notify_exception(None, "Error migrating SQL BACKEND toggle", {
                'domain': domain_name
            })
        return True
    else:
        return domain_object.use_sql_backend

def _should_use_sql_backend_in_tests(domain_object_or_name):
    """The default return value is False unless the ``TESTS_SHOULD_USE_SQL_BACKEND`` setting
    has been set or a Domain object with the same name exists."""
    assert settings.UNIT_TESTING
    from corehq.apps.domain.models import Domain
    override = getattr(settings, 'TESTS_SHOULD_USE_SQL_BACKEND', None)
    if override is not None:
        return override

    if domain_object_or_name and getattr(settings, 'DB_ENABLED', True):
        domain_object = domain_object_or_name \
            if isinstance(domain_object_or_name, Domain) \
            else Domain.get_by_name(domain_object_or_name)
        return domain_object and domain_object.use_sql_backend
    else:
        return False

def use_new_exports(domain_name):
    return NEW_EXPORTS.enabled(domain_name) or should_use_sql_backend(domain_name)

def use_sqlite_backend(domain_name):
    return TF_USES_SQLITE_BACKEND.enabled(domain_name) or should_use_sql_backend(domain_name)

def is_commcarecase(obj):
    from casexml.apps.case.models import CommCareCase
    from corehq.form_processor.models import CommCareCaseSQL
    return isinstance(obj, (CommCareCase, CommCareCaseSQL))
<|endoftext|>


<commit_before>import gevent
import logging
import time

from counterblock.lib import config
from counterblock.lib.processor import CaughtUpProcessor, CORE_FIRST_PRIORITY, CORE_LAST_PRIORITY, tasks, start_task

logger = logging.getLogger(__name__)

#TODO: Place any core CaughtUpProcessor tasks here<commit_msg>fix missing module<commit_after>import gevent
import logging
import time

from counterblock.lib import config
from counterblock.lib.processor import CaughtUpProcessor, CORE_FIRST_PRIORITY, CORE_LAST_PRIORITY, start_task

logger = logging.getLogger(__name__)

#TODO: Place any core CaughtUpProcessor tasks here<|endoftext|>