Source code for schrodinger.test.fixtures

"""
Test fixtures

Copyright Schrodinger, LLC. All rights reserved.
"""

import contextlib
import os


@contextlib.contextmanager
def use_static_opls_directory():
    """
    Sets OPLS_DIR to static data directory used in testing
    """

    # Avoid circular import
    from schrodinger.test import decorate_fflag
    with decorate_fflag.enable_feature_flag("OPLS_STATIC_PARAMETERS"):
        yield


class safe_patch_os_environ(contextlib.ContextDecorator):
    """
    Patch os.environ without clearing it to avoid crashes in threads

    Doesn't overwrite environment variables that are expected/allowed to change
    under pytest.
    """

    def __init__(self, values: dict = None):
        self._orig_env = {}
        self._orig_keys = set()
        self._values = values

    def __enter__(self):
        self.start()
        return self

    def __exit__(self, *args, **kwargs):
        self.stop()

    def start(self):
        orig_env = os.environ.copy()
        self._orig_keys = set(orig_env)
        ignore = ('PYTEST_CURRENT_TEST', 'SCHRODINGER_MMPREF_SANDBOX')
        for key in ignore:
            orig_env.pop(key, None)
        self._orig_env = orig_env
        if self._values:
            os.environ.update(self._values)

    def stop(self):
        # empty os.environ causes thread crashes (see MAE-46443)
        # mock.patch.dict temporarily clears os.environ when unpatching,
        # so we restore the original env without calling clear.
        new_keys = set(os.environ) - self._orig_keys
        for key in new_keys:
            del os.environ[key]
        os.environ.update(self._orig_env)