merge

Author: alexander-alderman-webb

tests/integrations/stdlib/test_httplib.py

@@ -202,12 +202,20 @@ def test_httplib_misuse(sentry_init, capture_events, request):
     )
 
 
-def test_outgoing_trace_headers(sentry_init, capture_events, monkeypatch):
-    # HTTPSConnection.send is passed a string containing (among other things)
-    # the headers on the request. Mock it so we can check the headers, and also
-    # so it doesn't try to actually talk to the internet.
-    mock_send = mock.Mock()
-    monkeypatch.setattr(HTTPSConnection, "send", mock_send)
+def test_outgoing_trace_headers(sentry_init, capture_events):
+    original_send = HTTPConnection.send
+
+    request_headers = {}
+
+    class HttpConnectionRecordingRequestHeaders(HTTPConnection):
+        def send(self, *args, **kwargs) -> None:
+            request_str = args[0]
+            for line in request_str.decode("utf-8").split("\r\n")[1:]:
+                if line:
+                    key, val = line.split(": ")
+                    request_headers[key] = val
+
+            original_send(self, *args, **kwargs)
 
     sentry_init(traces_sample_rate=1.0)
     events = capture_events()
@@ -229,20 +237,15 @@ def test_outgoing_trace_headers(sentry_init, capture_events, monkeypatch):
         op="greeting.sniff",
         trace_id="12312012123120121231201212312012",
     ) as transaction:
-        HTTPSConnection("www.squirrelchasers.com").request("GET", "/top-chasers")
-
-        (request_str,) = mock_send.call_args[0]
-        request_headers = {}
-        for line in request_str.decode("utf-8").split("\r\n")[1:]:
-            if line:
-                key, val = line.split(": ")
-                request_headers[key] = val
+        connection = HttpConnectionRecordingRequestHeaders("localhost", port=PORT)
+        connection.request("GET", "/top-chasers")
+        connection.getresponse()
 
     (event,) = events
     request_span = event["spans"][-1]
     expected_sentry_trace = "{trace_id}-{parent_span_id}-{sampled}".format(
-        trace_id=transaction.trace_id,
-        parent_span_id=request_span.span_id,
+        trace_id=event["contexts"]["trace"]["trace_id"],
+        parent_span_id=request_span["span_id"],
         sampled=1,
     )
     assert request_headers["sentry-trace"] == expected_sentry_trace
@@ -258,12 +261,20 @@ def test_outgoing_trace_headers(sentry_init, capture_events, monkeypatch):
     assert request_headers["baggage"] == expected_outgoing_baggage
 
 
-def test_outgoing_trace_headers_head_sdk(sentry_init, capture_events, monkeypatch):
-    # HTTPSConnection.send is passed a string containing (among other things)
-    # the headers on the request. Mock it so we can check the headers, and also
-    # so it doesn't try to actually talk to the internet.
-    mock_send = mock.Mock()
-    monkeypatch.setattr(HTTPSConnection, "send", mock_send)
+def test_outgoing_trace_headers_head_sdk(sentry_init, capture_events):
+    original_send = HTTPConnection.send
+
+    request_headers = {}
+
+    class HttpConnectionRecordingRequestHeaders(HTTPConnection):
+        def send(self, *args, **kwargs) -> None:
+            request_str = args[0]
+            for line in request_str.decode("utf-8").split("\r\n")[1:]:
+                if line:
+                    key, val = line.split(": ")
+                    request_headers[key] = val
+
+            original_send(self, *args, **kwargs)
 
     sentry_init(traces_sample_rate=0.5, release="foo")
     events = capture_events()
@@ -272,23 +283,18 @@ def test_outgoing_trace_headers_head_sdk(sentry_init, capture_events, monkeypatc
         transaction = continue_trace({})
 
     with start_transaction(transaction=transaction, name="Head SDK tx") as transaction:
-        connection = HTTPSConnection("www.squirrelchasers.com")
+        connection = HttpConnectionRecordingRequestHeaders("localhost", port=PORT)
         connection.request("GET", "/top-chasers")
-
-        (request_str,) = mock_send.call_args[0]
-        request_headers = {}
-        for line in request_str.decode("utf-8").split("\r\n")[1:]:
-            if line:
-                key, val = line.split(": ")
-                request_headers[key] = val
+        connection.getresponse()
 
     (event,) = events
-    request_span = event["spans"][0]
+    request_span = event["spans"][-1]
     expected_sentry_trace = "{trace_id}-{parent_span_id}-{sampled}".format(
-        trace_id=transaction.trace_id,
-        parent_span_id=request_span.span_id,
+        trace_id=event["contexts"]["trace"]["trace_id"],
+        parent_span_id=request_span["span_id"],
         sampled=1,
     )
+
     assert request_headers["sentry-trace"] == expected_sentry_trace
 
     expected_outgoing_baggage = (