Skip to content

Commit 26fc479

Browse files
committed
httplib patching erase headers
1 parent 9b62ced commit 26fc479

File tree

2 files changed

+51
-9
lines changed

2 files changed

+51
-9
lines changed

datadog_lambda/patch.py

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,10 @@
1717

1818
if sys.version_info >= (3, 0, 0):
1919
httplib_module = "http.client"
20+
from collections.abc import MutableMapping
2021
else:
2122
httplib_module = "httplib"
23+
from collections import MutableMapping
2224

2325
_httplib_patched = False
2426
_requests_patched = False
@@ -80,9 +82,9 @@ def _wrap_requests_request(func, instance, args, kwargs):
8082
into the outgoing requests.
8183
"""
8284
context = get_dd_trace_context()
83-
if "headers" in kwargs and isinstance(kwargs["headers"], dict):
85+
if "headers" in kwargs and isinstance(kwargs["headers"], MutableMapping):
8486
kwargs["headers"].update(context)
85-
elif len(args) >= 5 and isinstance(args[4], dict):
87+
elif len(args) >= 5 and isinstance(args[4], MutableMapping):
8688
args[4].update(context)
8789
else:
8890
kwargs["headers"] = context
@@ -100,9 +102,9 @@ def _wrap_httplib_request(func, instance, args, kwargs):
100102
the Datadog trace headers into the outgoing requests.
101103
"""
102104
context = get_dd_trace_context()
103-
if "headers" in kwargs and isinstance(kwargs["headers"], dict):
105+
if "headers" in kwargs and isinstance(kwargs["headers"], MutableMapping):
104106
kwargs["headers"].update(context)
105-
elif len(args) >= 4 and isinstance(args[3], dict):
107+
elif len(args) >= 4 and isinstance(args[3], MutableMapping):
106108
args[3].update(context)
107109
else:
108110
kwargs["headers"] = context

tests/test_patch.py

Lines changed: 45 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,9 @@
22
import unittest
33

44
try:
5-
from unittest.mock import patch
5+
from unittest.mock import patch, MagicMock
66
except ImportError:
7-
from mock import patch
7+
from mock import patch, MagicMock
88

99
from datadog_lambda.patch import _patch_httplib, _ensure_patch_requests
1010
from datadog_lambda.constants import TraceHeader
@@ -24,11 +24,51 @@ def setUp(self):
2424
def test_patch_httplib(self):
2525
_patch_httplib()
2626
if sys.version_info >= (3, 0, 0):
27-
import urllib.request as urllib
27+
from http.client import HTTPSConnection
2828
else:
29-
import urllib2 as urllib
30-
urllib.urlopen("https://www.datadoghq.com/")
29+
from httplib import HTTPSConnection
30+
31+
conn = HTTPSConnection("www.datadoghq.com")
32+
conn.request("GET", "/")
33+
conn.getresponse()
34+
35+
self.mock_get_dd_trace_context.assert_called()
36+
37+
def test_patch_httplib_dict_headers(self):
38+
_patch_httplib()
39+
if sys.version_info >= (3, 0, 0):
40+
from http.client import HTTPSConnection
41+
else:
42+
from httplib import HTTPSConnection
43+
44+
headers = MagicMock(spec=dict)
45+
headers["fake-header"] = "fake-value"
46+
47+
conn = HTTPSConnection("www.datadoghq.com")
48+
conn.request("GET", "/", headers=headers)
49+
conn.getresponse()
50+
51+
self.mock_get_dd_trace_context.assert_called()
52+
headers.update.assert_called()
53+
54+
def test_patch_httplib_dict_like_headers(self):
55+
_patch_httplib()
56+
if sys.version_info >= (3, 0, 0):
57+
from http.client import HTTPSConnection
58+
from collections.abc import MutableMapping
59+
else:
60+
from httplib import HTTPSConnection
61+
from collections import MutableMapping
62+
63+
headers = MagicMock(spec=MutableMapping)
64+
headers["fake-header"] = "fake-value"
65+
66+
conn = HTTPSConnection("www.datadoghq.com")
67+
conn.request("GET", "/", headers=headers)
68+
conn.getresponse()
69+
3170
self.mock_get_dd_trace_context.assert_called()
71+
headers.update.assert_called()
3272

3373
def test_patch_requests(self):
3474
_ensure_patch_requests()

0 commit comments

Comments
 (0)