test: add marker for running tests in subprocess (backport #3383) (#4403)

* test: add marker for running tests in subprocess (#3383)

The `subprocess` pytest marker can be used to run arbitrary Python code in a Python
subprocess. This is meant to replace the existing fixture to allow actual
Python code to be written and tested, instead of using literal strings.

## Checklist
- [ ] Added to the correct milestone.
- [ ] Tests provided or description of manual testing performed is included in the code or PR.
- [ ] Library documentation is updated.
- [ ] [Corp site](https://github.com/DataDog/documentation/) documentation is updated (link to the PR).

(cherry picked from commit c437ad2)

# Conflicts:
#	tests/contrib/httpx/test_httpx.py
#	tests/integration/test_integration.py
#	tests/tracer/test_encoders.py

* fix conflicts

* Use pytest marker to generate snapshot, set variants

pull/4418

Co-authored-by: Gabriele N. Tornetta <[email protected]>
Co-authored-by: Munir Abdinur <[email protected]>
Co-authored-by: Munir Abdinur <[email protected]>

Author: 3 people

conftest.py

@@ -28,6 +28,25 @@
 # Hook for dynamic configuration of pytest in CI
 # https://docs.pytest.org/en/6.2.1/reference.html#pytest.hookspec.pytest_configure
 def pytest_configure(config):
+    config.addinivalue_line(
+        "markers",
+        """subprocess(status, out, err, args, env, parametrize, ddtrace_run):
+            Mark test functions whose body is to be run as stand-alone Python
+            code in a subprocess.
+
+            Arguments:
+                status: the expected exit code of the subprocess.
+                out: the expected stdout of the subprocess, or None to ignore.
+                err: the expected stderr of the subprocess, or None to ignore.
+                args: the command line arguments to pass to the subprocess.
+                env: the environment variables to override for the subprocess.
+                parametrize: whether to parametrize the test function. This is
+                    similar to the `parametrize` marker, but arguments are
+                    passed to the subprocess via environment variables.
+                ddtrace_run: whether to run the test using ddtrace-run.
+        """,
+    )
+
     if os.getenv("CI") != "true":
         return
 

tests/conftest.py

@@ -1,8 +1,15 @@
+import ast
 import contextlib
+from itertools import product
 import os
 import sys
+from tempfile import NamedTemporaryFile
+import time
 
+from _pytest.runner import CallInfo
+from _pytest.runner import TestReport
 import pytest
+from six import PY2
 
 from tests.utils import DummyTracer
 from tests.utils import TracerSpanContainer
@@ -94,3 +101,143 @@ def _snapshot(**kwargs):
             yield snapshot
 
     return _snapshot
+
+
+# DEV: The dump_code_to_file function is adapted from the compile function in
+# the py_compile module of the Python standard library. It generates .pyc files
+# with the right format.
+if PY2:
+    import marshal
+    from py_compile import MAGIC
+    from py_compile import wr_long
+
+    from _pytest._code.code import ExceptionInfo
+
+    def dump_code_to_file(code, file):
+        file.write(MAGIC)
+        wr_long(file, long(time.time()))  # noqa
+        marshal.dump(code, file)
+        file.flush()
+
+
+else:
+    import importlib
+
+    code_to_pyc = getattr(
+        importlib._bootstrap_external, "_code_to_bytecode" if sys.version_info < (3, 7) else "_code_to_timestamp_pyc"
+    )
+
+    def dump_code_to_file(code, file):
+        file.write(code_to_pyc(code, time.time(), len(code.co_code)))
+        file.flush()
+
+
+def unwind_params(params):
+    if params is None:
+        yield None
+        return
+
+    for _ in product(*(((k, v) for v in vs) for k, vs in params.items())):
+        yield dict(_)
+
+
+class FunctionDefFinder(ast.NodeVisitor):
+    def __init__(self, func_name):
+        super(FunctionDefFinder, self).__init__()
+        self.func_name = func_name
+        self._body = None
+
+    def generic_visit(self, node):
+        return self._body or super(FunctionDefFinder, self).generic_visit(node)
+
+    def visit_FunctionDef(self, node):
+        if node.name == self.func_name:
+            self._body = node.body
+
+    def find(self, file):
+        with open(file) as f:
+            t = ast.parse(f.read())
+            self.visit(t)
+            t.body = self._body
+            return t
+
+
+def run_function_from_file(item, params=None):
+    file, _, func = item.location
+    marker = item.get_closest_marker("subprocess")
+
+    file_index = 1
+    args = marker.kwargs.get("args", [])
+    args.insert(0, None)
+    args.insert(0, sys.executable)
+    if marker.kwargs.get("ddtrace_run", False):
+        file_index += 1
+        args.insert(0, "ddtrace-run")
+
+    env = os.environ.copy()
+    env.update(marker.kwargs.get("env", {}))
+    if params is not None:
+        env.update(params)
+
+    expected_status = marker.kwargs.get("status", 0)
+
+    expected_out = marker.kwargs.get("out", "")
+    if expected_out is not None:
+        expected_out = expected_out.encode("utf-8")
+
+    expected_err = marker.kwargs.get("err", "")
+    if expected_err is not None:
+        expected_err = expected_err.encode("utf-8")
+
+    with NamedTemporaryFile(mode="wb", suffix=".pyc") as fp:
+        dump_code_to_file(compile(FunctionDefFinder(func).find(file), file, "exec"), fp.file)
+
+        start = time.time()
+        args[file_index] = fp.name
+        out, err, status, _ = call_program(*args, env=env)
+        end = time.time()
+        excinfo = None
+
+        if status != expected_status:
+            excinfo = AssertionError(
+                "Expected status %s, got %s.\n=== Captured STDERR ===\n%s=== End of captured STDERR ==="
+                % (expected_status, status, err.decode("utf-8"))
+            )
+        elif expected_out is not None and out != expected_out:
+            excinfo = AssertionError("STDOUT: Expected [%s] got [%s]" % (expected_out, out))
+        elif expected_err is not None and err != expected_err:
+            excinfo = AssertionError("STDERR: Expected [%s] got [%s]" % (expected_err, err))
+
+        if PY2 and excinfo is not None:
+            try:
+                raise excinfo
+            except Exception:
+                excinfo = ExceptionInfo(sys.exc_info())
+
+        call_info_args = dict(result=None, excinfo=excinfo, start=start, stop=end, when="call")
+        if not PY2:
+            call_info_args["duration"] = end - start
+
+        return TestReport.from_item_and_call(item, CallInfo(**call_info_args))
+
+
+@pytest.hookimpl(tryfirst=True)
+def pytest_runtest_protocol(item):
+    marker = item.get_closest_marker("subprocess")
+    if marker:
+        params = marker.kwargs.get("parametrize", None)
+        ihook = item.ihook
+        base_name = item.nodeid
+
+        for ps in unwind_params(params):
+            nodeid = (base_name + str(ps)) if ps is not None else base_name
+
+            ihook.pytest_runtest_logstart(nodeid=nodeid, location=item.location)
+
+            report = run_function_from_file(item, ps)
+            report.nodeid = nodeid
+            ihook.pytest_runtest_logreport(report=report)
+
+            ihook.pytest_runtest_logfinish(nodeid=nodeid, location=item.location)
+
+        return True

tests/contrib/gevent/test_monkeypatch.py

@@ -20,20 +20,17 @@ def test_gevent_warning(monkeypatch):
     assert b"RuntimeWarning: Loading ddtrace before using gevent monkey patching" in subp.stderr.read()
 
 
-def test_gevent_auto_patching(run_python_code_in_subprocess):
-    code = """
-import ddtrace; ddtrace.patch_all()
-
-import gevent  # Patch on import
-from ddtrace.contrib.gevent import GeventContextProvider
+@pytest.mark.subprocess
+def test_gevent_auto_patching():
+    import ddtrace
 
+    ddtrace.patch_all()
+    # Patch on import
+    import gevent  # noqa
 
-assert isinstance(ddtrace.tracer.context_provider, GeventContextProvider)
-"""
+    from ddtrace.contrib.gevent import GeventContextProvider
 
-    out, err, status, pid = run_python_code_in_subprocess(code)
-    assert status == 0, err
-    assert out == b""
+    assert isinstance(ddtrace.tracer.context_provider, GeventContextProvider)
 
 
 def test_gevent_ddtrace_run_auto_patching(ddtrace_run_python_code_in_subprocess):