8000 httplib patching erase headers · DataDog/datadog-lambda-python@26fc479 · GitHub
[go: up one dir, main page]
Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Appearance settings

Commit 26fc479

Browse files
tianchutianchu
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
< 8000 td data-grid-cell-id="diff-73e003e723da48c6cdf2f12de86b60ca6a04761c7d94c8aa6d3b22001b712eac-103-104-0" data-selected="false" role="gridcell" style="background-color:var(--diffBlob-deletionNum-bgColor, var(--diffBlob-deletion-bgColor-num));text-align:center" tabindex="-1" valign="top" class="focusable-grid-cell diff-line-number position-relative left-side">103
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()
-
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)
0