From 20b44937f7c7b61fe5cac14b907d18747ac0ab50 Mon Sep 17 00:00:00 2001
From: "Tahir H. Butt" <tahir.butt@datadoghq.com>
Date: Wed, 21 Oct 2020 20:33:40 -0400
Subject: [PATCH 1/2] fix snapshot decorator for pytest fixtures

---
 tests/__init__.py | 93 +++++++++++++++++++++++------------------------
 1 file changed, 45 insertions(+), 48 deletions(-)

diff --git a/tests/__init__.py b/tests/__init__.py
index e3f60b98942..eba40c2e042 100644
--- a/tests/__init__.py
+++ b/tests/__init__.py
@@ -1,11 +1,11 @@
 import contextlib
-import functools
 import inspect
 import os
 import sys
 from contextlib import contextmanager
 
 import pytest
+from ddtrace.vendor import wrapt
 
 import ddtrace
 from ddtrace import Tracer, Span
@@ -815,54 +815,51 @@ def snapshot(ignores=None, tracer=ddtrace.tracer):
     """
     ignores = ignores or []
 
-    def dec(f):
-        @functools.wraps(f)
-        def wrapper(*args, **kwargs):
-            if len(args) > 1:
-                self = args[0]
-                clsname = self.__class__.__name__
-            else:
-                clsname = ""
+    @wrapt.decorator
+    def wrapper(wrapped, instance, args, kwargs):
+        if len(args) > 1:
+            self = args[0]
+            clsname = self.__class__.__name__
+        else:
+            clsname = ""
 
-            module = inspect.getmodule(f)
+        module = inspect.getmodule(wrapped)
 
-            # Use the fully qualified function name as a unique test token to
-            # identify the snapshot.
-            token = "{}{}{}.{}".format(module.__name__, "." if clsname else "", clsname, f.__name__)
+        # Use the fully qualified function name as a unique test token to
+        # identify the snapshot.
+        token = "{}{}{}.{}".format(module.__name__, "." if clsname else "", clsname, wrapped.__name__)
 
-            conn = httplib.HTTPConnection(tracer.writer.api.hostname, tracer.writer.api.port)
+        conn = httplib.HTTPConnection(tracer.writer.api.hostname, tracer.writer.api.port)
+        try:
+            # Signal the start of this test case to the test agent.
             try:
-                # Signal the start of this test case to the test agent.
-                try:
-                    conn.request("GET", "/test/start?token=%s" % token)
-                except Exception as e:
-                    pytest.fail("Could not connect to test agent: %s" % str(e), pytrace=False)
-
-                r = conn.getresponse()
-                if r.status != 200:
-                    # The test agent returns nice error messages we can forward to the user.
-                    raise SnapshotFailed(r.read().decode())
-
-                # Run the test.
-                ret = f(*args, **kwargs)
-
-                # Flush out any remnant traces.
-                tracer.writer.flush_queue()
-
-                # Query for the results of the test.
-                conn = httplib.HTTPConnection(tracer.writer.api.hostname, tracer.writer.api.port)
-                conn.request("GET", "/test/snapshot?ignores=%s&token=%s" % (",".join(ignores), token))
-                r = conn.getresponse()
-                if r.status != 200:
-                    raise SnapshotFailed(r.read().decode())
-                return ret
-            except SnapshotFailed as e:
-                # Fail the test if a failure has occurred and print out the
-                # message we got from the test agent.
-                pytest.fail(str(e), pytrace=False)
-            finally:
-                conn.close()
-
-        return wrapper
-
-    return dec
+                conn.request("GET", "/test/start?token=%s" % token)
+            except Exception as e:
+                pytest.fail("Could not connect to test agent: %s" % str(e), pytrace=False)
+
+            r = conn.getresponse()
+            if r.status != 200:
+                # The test agent returns nice error messages we can forward to the user.
+                raise SnapshotFailed(r.read().decode())
+
+            # Run the test.
+            ret = wrapped(*args, **kwargs)
+
+            # Flush out any remnant traces.
+            tracer.writer.flush_queue()
+
+            # Query for the results of the test.
+            conn = httplib.HTTPConnection(tracer.writer.api.hostname, tracer.writer.api.port)
+            conn.request("GET", "/test/snapshot?ignores=%s&token=%s" % (",".join(ignores), token))
+            response = conn.getresponse()
+            if response.status != 200:
+                raise SnapshotFailed(response.read().decode())
+            return ret
+        except SnapshotFailed as e:
+            # Fail the test if a failure has occurred and print out the
+            # message we got from the test agent.
+            pytest.fail(str(e), pytrace=False)
+        finally:
+            conn.close()
+
+    return wrapper

From 6a0731495c4a4bfcb65fdf5e335b6e1d303a25fd Mon Sep 17 00:00:00 2001
From: "Tahir H. Butt" <tahir.butt@datadoghq.com>
Date: Thu, 22 Oct 2020 14:23:54 -0400
Subject: [PATCH 2/2] flush before starting snapshot

---
 tests/__init__.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tests/__init__.py b/tests/__init__.py
index eba40c2e042..89757080de7 100644
--- a/tests/__init__.py
+++ b/tests/__init__.py
@@ -831,6 +831,10 @@ def wrapper(wrapped, instance, args, kwargs):
 
         conn = httplib.HTTPConnection(tracer.writer.api.hostname, tracer.writer.api.port)
         try:
+            # clear queue in case traces have been generated before test case is
+            # itself run
+            tracer.writer.flush_queue()
+
             # Signal the start of this test case to the test agent.
             try:
                 conn.request("GET", "/test/start?token=%s" % token)