fix(ci): wait for traces to be sent in snapshot tests (#4374) (#4384)

Add a wait_for_traces argument to snapshots to poll the test agent until the correct number of traces are available. In particular httpx has been flaky.

(cherry picked from commit 4e710ca)

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

Author: 3 people

tests/contrib/httpx/test_httpx.py

@@ -169,13 +169,13 @@ def test_configure_service_name_env():
 
     async def test():
         token = "tests.contrib.httpx.test_httpx.test_configure_service_name_env"
-        with snapshot_context(token=token):
+        with snapshot_context(wait_for_num_traces=1, token=token):
             DEFAULT_HEADERS = {
                 "User-Agent": "python-httpx/x.xx.x",
             }
             httpx.get(url, headers=DEFAULT_HEADERS)
 
-        with snapshot_context(token=token):
+        with snapshot_context(wait_for_num_traces=1, token=token):
             async with httpx.AsyncClient() as client:
                 DEFAULT_HEADERS = {
                     "User-Agent": "python-httpx/x.xx.x",
@@ -209,13 +209,13 @@ def test_configure_global_service_name_env():
 
     async def test():
         token = "tests.contrib.httpx.test_httpx.test_configure_global_service_name_env"
-        with snapshot_context(token=token):
+        with snapshot_context(wait_for_num_traces=1, token=token):
             DEFAULT_HEADERS = {
                 "User-Agent": "python-httpx/x.xx.x",
             }
             httpx.get(url, headers=DEFAULT_HEADERS)
 
-        with snapshot_context(token=token):
+        with snapshot_context(wait_for_num_traces=1, token=token):
             async with httpx.AsyncClient() as client:
                 await client.get(url, headers=DEFAULT_HEADERS)
 
@@ -232,11 +232,11 @@ async def test_get_500(snapshot_context):
         We mark the span as an error
     """
     url = get_url("/status/500")
-    with snapshot_context():
+    with snapshot_context(wait_for_num_traces=1):
         resp = httpx.get(url, headers=DEFAULT_HEADERS)
         assert resp.status_code == 500
 
-    with snapshot_context():
+    with snapshot_context(wait_for_num_traces=1):
         async with httpx.AsyncClient() as client:
             resp = await client.get(url, headers=DEFAULT_HEADERS)
             assert resp.status_code == 500
@@ -245,17 +245,17 @@ async def test_get_500(snapshot_context):
 @pytest.mark.asyncio
 async def test_split_by_domain(snapshot_context):
     """
-    When split_by_domain is configure
+    When split_by_domain is configured
         We set the service name to the <host>:<port>
     """
     url = get_url("/status/200")
 
     with override_config("httpx", {"split_by_domain": True}):
-        with snapshot_context():
+        with snapshot_context(wait_for_num_traces=1):
             resp = httpx.get(url, headers=DEFAULT_HEADERS)
             assert resp.status_code == 200
 
-        with snapshot_context():
+        with snapshot_context(wait_for_num_traces=1):
             async with httpx.AsyncClient() as client:
                 resp = await client.get(url, headers=DEFAULT_HEADERS)
                 assert resp.status_code == 200
@@ -275,11 +275,11 @@ async def test_trace_query_string(snapshot_context):
         "User-Agent": "python-httpx/x.xx.x",
     }
     with override_http_config("httpx", {"trace_query_string": True}):
-        with snapshot_context():
+        with snapshot_context(wait_for_num_traces=1):
             resp = httpx.get(url, headers=headers)
             assert resp.status_code == 200
 
-        with snapshot_context():
+        with snapshot_context(wait_for_num_traces=1):
             async with httpx.AsyncClient() as client:
                 resp = await client.get(url, headers=headers)
                 assert resp.status_code == 200
@@ -300,11 +300,11 @@ async def test_request_headers(snapshot_context):
 
     try:
         config.httpx.http.trace_headers(["Some-Request-Header", "Some-Response-Header"])
-        with snapshot_context():
+        with snapshot_context(wait_for_num_traces=1):
             resp = httpx.get(url, headers=headers)
             assert resp.status_code == 200
 
-        with snapshot_context():
+        with snapshot_context(wait_for_num_traces=1):
             async with httpx.AsyncClient() as client:
                 resp = await client.get(url, headers=headers)
                 assert resp.status_code == 200

tests/utils.py

@@ -1,9 +1,11 @@
 import contextlib
 from contextlib import contextmanager
 import inspect
+import json
 import os
 import subprocess
 import sys
+import time
 from typing import List
 
 import attr
@@ -854,7 +856,7 @@ def clear(self):
 
 
 @contextmanager
-def snapshot_context(token, ignores=None, tracer=None, async_mode=True, variants=None):
+def snapshot_context(token, ignores=None, tracer=None, async_mode=True, variants=None, wait_for_num_traces=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:
@@ -910,6 +912,27 @@ def snapshot_context(token, ignores=None, tracer=None, async_mode=True, variants
                 del tracer._writer._headers["X-Datadog-Test-Session-Token"]
                 del os.environ["_DD_TRACE_WRITER_ADDITIONAL_HEADERS"]
 
+        conn = httplib.HTTPConnection(parsed.hostname, parsed.port)
+
+        # Wait for the traces to be available
+        if wait_for_num_traces is not None:
+            traces = []
+            for i in range(20):
+                try:
+                    conn.request("GET", "/test/session/traces?test_session_token=%s" % token)
+                    r = conn.getresponse()
+                    if r.status == 200:
+                        traces = json.loads(r.read())
+                        if len(traces) == wait_for_num_traces:
+                            break
+                except Exception:
+                    pass
+                time.sleep(0.1)
+            else:
+                pytest.fail(
+                    "Expected %r trace(s), got %r:\n%s" % (wait_for_num_traces, len(traces), traces), pytrace=False
+                )
+
         # Query for the results of the test.
         conn = httplib.HTTPConnection(parsed.hostname, parsed.port)
         conn.request("GET", "/test/session/snapshot?ignores=%s&test_session_token=%s" % (",".join(ignores), token))