Add pytest snapshot fixture (#2688)

The snapshot decorator worked ok for basic integration tests but isn't
compatible for more advanced pytest usage like parametrization.

By creating a fixture we can `pytest.mark.snapshot` test cases which
will make use of a `snapshot` pytest fixture which supports more
advanced pytest usage.

Co-authored-by: Brett Langdon <[email protected]>

Author: Kyle-Verhoog

tests/conftest.py

@@ -5,6 +5,13 @@
 from tests.utils import DummyTracer
 from tests.utils import TracerSpanContainer
 from tests.utils import call_program
+from tests.utils import snapshot_context
+
+
+def pytest_configure(config):
+    config.addinivalue_line(
+        "markers", "snapshot(*args, **kwargs): mark test to run as a snapshot test which sends traces to the test agent"
+    )
 
 
 @pytest.fixture
@@ -37,3 +44,19 @@ def _run(code, **kwargs):
         return call_program("ddtrace-run", sys.executable, str(pyfile), **kwargs)
 
     yield _run
+
+
+@pytest.fixture(autouse=True)
+def snapshot(request):
+    marks = [m for m in request.node.iter_markers(name="snapshot")]
+    assert len(marks) < 2, "Multiple snapshot marks detected"
+    if marks:
+        snap = marks[0]
+        token = ""
+        token += request.module.__name__
+        token += ".%s" % request.cls.__name__ if request.cls else ""
+        token += ".%s" % request.node.name
+        with snapshot_context(token, *snap.args, **snap.kwargs) as snapshot:
+            yield snapshot
+    else:
+        yield

tests/utils.py

@@ -6,6 +6,7 @@
 import sys
 from typing import List
 
+import attr
 import pytest
 
 import ddtrace
@@ -19,6 +20,7 @@
 from ddtrace.internal.compat import to_unicode
 from ddtrace.internal.encoding import JSONEncoder
 from ddtrace.internal.writer import AgentWriter
+from ddtrace.utils.formats import parse_tags_str
 from ddtrace.vendor import wrapt
 from tests.subprocesstest import SubprocessTestCase
 
@@ -820,8 +822,30 @@ class SnapshotFailed(Exception):
     pass
 
 
+@attr.s
+class SnapshotTest(object):
+    token = attr.ib(type=str)
+    tracer = attr.ib(type=ddtrace.Tracer, default=ddtrace.tracer)
+
+    def clear(self):
+        """Clear any traces sent that were sent for this snapshot."""
+        parsed = parse.urlparse(self.tracer.writer.agent_url)
+        conn = httplib.HTTPConnection(parsed.hostname, parsed.port)
+        conn.request("GET", "/test/clear?token=%s" % self.token)
+        resp = conn.getresponse()
+        assert resp.status == 200
+
+
 @contextmanager
-def snapshot_context(token, ignores=None, tracer=None, async_mode=True):
+def snapshot_context(token, ignores=None, tracer=None, async_mode=True, variants=None):
+    # Use variant that applies to update test token. One must apply. If none
+    # apply, the test should have been marked as skipped.
+    if variants:
+        applicable_variant_ids = [k for (k, v) in variants.items() if v]
+        assert len(applicable_variant_ids) == 1
+        variant_id = applicable_variant_ids[0]
+        token = "{}_{}".format(token, variant_id) if variant_id else token
+
     ignores = ignores or []
     if not tracer:
         tracer = ddtrace.tracer
@@ -839,6 +863,13 @@ def snapshot_context(token, ignores=None, tracer=None, async_mode=True):
         if async_mode:
             # Patch the tracer writer to include the test token header for all requests.
             tracer.writer._headers["X-Datadog-Test-Token"] = token
+
+            # Also add a header to the environment for subprocesses test cases that might use snapshotting.
+            existing_headers = parse_tags_str(os.environ.get("_DD_TRACE_WRITER_ADDITIONAL_HEADERS", ""))
+            existing_headers.update({"X-Datadog-Test-Token": token})
+            os.environ["_DD_TRACE_WRITER_ADDITIONAL_HEADERS"] = ",".join(
+                ["%s:%s" % (k, v) for k, v in existing_headers.items()]
+            )
         else:
             # Signal the start of this test case to the test agent.
             try:
@@ -851,14 +882,17 @@ def snapshot_context(token, ignores=None, tracer=None, async_mode=True):
                     # The test agent returns nice error messages we can forward to the user.
                     raise SnapshotFailed(r.read())
 
-        # Return context to the caller
         try:
-            yield
+            yield SnapshotTest(
+                tracer=tracer,
+                token=token,
+            )
         finally:
             # Force a flush so all traces are submitted.
             tracer.writer.flush_queue()
             if async_mode:
                 del tracer.writer._headers["X-Datadog-Test-Token"]
+                del os.environ["_DD_TRACE_WRITER_ADDITIONAL_HEADERS"]
 
         # Query for the results of the test.
         conn = httplib.HTTPConnection(parsed.hostname, parsed.port)
@@ -917,15 +951,7 @@ def wrapper(wrapped, instance, args, kwargs):
             else token_override
         )
 
-        # Use variant that applies to update test token. One must apply. If none
-        # apply, the test should have been marked as skipped.
-        if variants:
-            applicable_variant_ids = [k for (k, v) in variants.items() if v]
-            assert len(applicable_variant_ids) == 1
-            variant_id = applicable_variant_ids[0]
-            token = "{}_{}".format(token, variant_id) if variant_id else token
-
-        with snapshot_context(token, ignores=ignores, tracer=tracer, async_mode=async_mode):
+        with snapshot_context(token, ignores=ignores, tracer=tracer, async_mode=async_mode, variants=variants):
             # Run the test.
             if include_tracer:
                 kwargs["tracer"] = tracer