Source code for schrodinger.test.decorate_fflag

"""
Enable or disable a feature flag in a specific situation.  Can be used as a
decorator for a TestCase class or a test function, or as a context manager.
See nose_plugin_test.py for examples. Multiple feature flag toggling
decorators/contexts can be nested to define desired feature flag environment.

@copyright: (c) Schrodinger, LLC All rights reserved.
"""

import os
from contextlib import ContextDecorator
from unittest.mock import patch

from schrodinger.utils import mmutil


class enable_feature_flag(ContextDecorator):
    """
    Provides a decorator/context manager for enabling feature flags.  Methods
    are named to allow this to be used the same ways as Mock's patches:

    * As context manager
    * As function or test class decorator
    * Explicit start and stop.
    """
    on = 1

    def __init__(self, flag):
        self.flag = flag
        self.original_state = None

    def start(self):
        """
        Starts the patching.  Used both in the decorator and the context
        manager.
        """
        current_environ = os.environ.get('SCHRODINGER_FEATURE_FLAGS', '')
        desired_environ = self._set_ff(self.flag, current_environ, self.on)

        # Some job control fflags must be set in this special variable to be
        # accessible from perl.
        perl_current_environ = os.environ.get(
            'SCHRODINGER_PERL_FEATURE_FLAGS_ENABLED', '')
        if self.flag not in perl_current_environ:
            perl_desired_environ = perl_current_environ + ' ' + self.flag
            perl_desired_environ.strip()
        else:
            perl_desired_environ = perl_current_environ

        self.environ_patcher = patch.dict(
            'os.environ', {
                'SCHRODINGER_FEATURE_FLAGS': desired_environ,
                'SCHRODINGER_PERL_FEATURE_FLAGS_ENABLED': perl_desired_environ
            })

        self.original_state = mmutil.feature_flag_is_enabled_s(self.flag)
        mmutil.feature_flags_set_user_state_s(self.flag, self.on)
        self.environ_patcher.start()

    def stop(self):
        """
        Stops the patching.  Used both in the decorator context and the
        context manager context.
        """
        self.environ_patcher.stop()
        mmutil.feature_flags_set_user_state_s(self.flag, self.original_state)

    def _onstr(self, on=False):
        """Return the +/- representation of a boolean value."""
        if on:
            return '+'
        else:
            return '-'

    def _set_ff(self, flagname, flagstr, on=True):
        """Create the featureflag environment variable string."""
        flags = set(flagstr.split())
        desired = self._onstr(on) + flagname
        opposite = self._onstr(not on) + flagname
        if opposite in flags:
            flags = flags - {opposite}
        flags.add(desired)
        return ' '.join(flags)

    def __enter__(self):
        """Enter the context manager."""
        self.start()

    def __exit__(self, exc_type, exc_value, traceback):
        """Exit the context manager."""
        self.stop()


class disable_feature_flag(enable_feature_flag):
    """
    Disable a feature flag as either a context manager or a decorator for a
    unittest.TestCase or a test function.
    :type flag: str
    :param flag: Name of the feature flag to be disabled.
    """
    on = 0