1"""
2subprocesstest enables unittest test cases and suites to be run in separate
3python interpreter instances.
4
5A base class SubprocessTestCase is provided that, when extended, will run test
6cases marked with @run_in_subprocess in a separate python interpreter.
7"""
8import inspect
9import os
10import subprocess
11import sys
12import unittest
13
14
15SUBPROC_TEST_ATTR = "_subproc_test"
16SUBPROC_TEST_ENV_ATTR = "_subproc_test_env"
17SUBPROC_ENV_VAR = "SUBPROCESS_TEST"
18
19
20def run_in_subprocess(*args, **kwargs):
21    """
22    Marks a test case that is to be run in its own 'clean' interpreter instance.
23
24    When applied to a SubprocessTestCase class, each method will be run in a
25    separate interpreter instance.
26
27    When applied only to a method of a SubprocessTestCase, only the method
28    will be run in a separate interpreter instance.
29
30    Usage on a class::
31
32        from tests.subprocesstest import SubprocessTestCase, run_in_subprocess
33
34        @run_in_subprocess
35        class PatchTests(SubprocessTestCase):
36            # will be run in new interpreter
37            def test_patch_before_import(self):
38                patch()
39                import module
40
41            # will be run in new interpreter as well
42            def test_patch_after_import(self):
43                import module
44                patch()
45
46
47    Usage on a test method::
48
49        class OtherTests(SubprocessTestCase):
50            @run_in_subprocess
51            def test_case(self):
52                # Run in subprocess
53                pass
54
55            def test_case(self):
56                # NOT run in subprocess
57                pass
58
59    :param env_override: dict of environment variables to provide to the subprocess.
60    :return:
61    """
62    env_overrides = kwargs.get("env_overrides")
63
64    def wrapper(obj):
65        setattr(obj, SUBPROC_TEST_ATTR, True)
66        if env_overrides is not None:
67            setattr(obj, SUBPROC_TEST_ENV_ATTR, env_overrides)
68        return obj
69
70    # Support both @run_in_subprocess and @run_in_subprocess(env_overrides=...) usage
71    if len(args) == 1 and callable(args[0]):
72        return wrapper(args[0])
73    else:
74        return wrapper
75
76
77class SubprocessTestCase(unittest.TestCase):
78    run_in_subprocess = staticmethod(run_in_subprocess)
79
80    def _full_method_name(self):
81        test = getattr(self, self._testMethodName)
82        # DEV: we have to use the internal self reference of the bound test
83        # method to pull out the class and module since using a mix of `self`
84        # and the test attributes will result in inconsistencies when the test
85        # method is defined on another class.
86        # A concrete case of this is a parent and child TestCase where the child
87        # doesn't override a parent test method. The full_method_name we want
88        # is that of the child test method (even though it exists on the parent).
89        # This is only true if the test method is bound by pytest; pytest>=5.4 returns a function.
90        if inspect.ismethod(test):
91            modpath = test.__self__.__class__.__module__
92            clsname = test.__self__.__class__.__name__
93        else:
94            modpath = self.__class__.__module__
95            clsname = self.__class__.__name__
96        testname = test.__name__
97        testcase_name = "{}.{}.{}".format(modpath, clsname, testname)
98        return testcase_name
99
100    def _run_test_in_subprocess(self, result):
101        full_testcase_name = self._full_method_name()
102
103        # Copy the environment and include the special subprocess environment
104        # variable for the subprocess to detect.
105        env_overrides = self._get_env_overrides()
106        sp_test_env = os.environ.copy()
107        sp_test_env.update(env_overrides)
108        sp_test_env[SUBPROC_ENV_VAR] = "True"
109        sp_test_cmd = ["python", "-m", "unittest", full_testcase_name]
110        sp = subprocess.Popen(
111            sp_test_cmd,
112            stdout=subprocess.PIPE,
113            stderr=subprocess.PIPE,
114            env=sp_test_env,
115        )
116        stdout, stderr = sp.communicate()
117
118        if sp.returncode:
119            try:
120                cmdf = " ".join(sp_test_cmd)
121                raise Exception('Subprocess Test "{}" Failed'.format(cmdf))
122            except Exception:
123                exc_info = sys.exc_info()
124
125            # DEV: stderr, stdout are byte sequences so to print them nicely
126            #      back out they should be decoded.
127            sys.stderr.write(stderr.decode())
128            sys.stdout.write(stdout.decode())
129            result.addFailure(self, exc_info)
130        else:
131            result.addSuccess(self)
132
133    def _in_subprocess(self):
134        """Determines if the test is being run in a subprocess.
135
136        This is done by checking for an environment variable that we call the
137        subprocess test with.
138
139        :return: whether the test is a subprocess test
140        """
141        return os.getenv(SUBPROC_ENV_VAR, None) is not None
142
143    def _is_subprocess_test(self):
144        if hasattr(self, SUBPROC_TEST_ATTR):
145            return True
146
147        test = getattr(self, self._testMethodName)
148        if hasattr(test, SUBPROC_TEST_ATTR):
149            return True
150
151        return False
152
153    def _get_env_overrides(self):
154        if hasattr(self, SUBPROC_TEST_ENV_ATTR):
155            return getattr(self, SUBPROC_TEST_ENV_ATTR)
156
157        test = getattr(self, self._testMethodName)
158        if hasattr(test, SUBPROC_TEST_ENV_ATTR):
159            return getattr(test, SUBPROC_TEST_ENV_ATTR)
160
161        return {}
162
163    def run(self, result=None):
164        if not self._is_subprocess_test():
165            return super(SubprocessTestCase, self).run(result=result)
166
167        if self._in_subprocess():
168            return super(SubprocessTestCase, self).run(result=result)
169        else:
170            self._run_test_in_subprocess(result)
171