httplib patching erase headers

Author: tianchu

datadog_lambda/patch.py

@@ -17,8 +17,10 @@
 
 if sys.version_info >= (3, 0, 0):
     httplib_module = "http.client"
+    from collections.abc import MutableMapping
 else:
     httplib_module = "httplib"
+    from collections import MutableMapping
 
 _httplib_patched = False
 _requests_patched = False
@@ -80,9 +82,9 @@ def _wrap_requests_request(func, instance, args, kwargs):
     into the outgoing requests.
     """
     context = get_dd_trace_context()
-    if "headers" in kwargs and isinstance(kwargs["headers"], dict):
+    if "headers" in kwargs and isinstance(kwargs["headers"], MutableMapping):
         kwargs["headers"].update(context)
-    elif len(args) >= 5 and isinstance(args[4], dict):
+    elif len(args) >= 5 and isinstance(args[4], MutableMapping):
         args[4].update(context)
     else:
         kwargs["headers"] = context
@@ -100,9 +102,9 @@ def _wrap_httplib_request(func, instance, args, kwargs):
     the Datadog trace headers into the outgoing requests.
     """
     context = get_dd_trace_context()
-    if "headers" in kwargs and isinstance(kwargs["headers"], dict):
+    if "headers" in kwargs and isinstance(kwargs["headers"], MutableMapping):
         kwargs["headers"].update(context)
-    elif len(args) >= 4 and isinstance(args[3], dict):
+    elif len(args) >= 4 and isinstance(args[3], MutableMapping):
         args[3].update(context)
     else:
         kwargs["headers"] = context

tests/test_patch.py

@@ -2,9 +2,9 @@
 import unittest
 
 try:
-    from unittest.mock import patch
+    from unittest.mock import patch, MagicMock
 except ImportError:
-    from mock import patch
+    from mock import patch, MagicMock
 
 from datadog_lambda.patch import _patch_httplib, _ensure_patch_requests
 from datadog_lambda.constants import TraceHeader
@@ -24,11 +24,51 @@ def setUp(self):
     def test_patch_httplib(self):
         _patch_httplib()
         if sys.version_info >= (3, 0, 0):
-            import urllib.request as urllib
+            from http.client import HTTPSConnection
         else:
-            import urllib2 as urllib
-        urllib.urlopen("https://www.datadoghq.com/")
+            from httplib import HTTPSConnection
+
+        conn = HTTPSConnection("www.datadoghq.com")
+        conn.request("GET", "/")
+        conn.getresponse()
+
+        self.mock_get_dd_trace_context.assert_called()
+
+    def test_patch_httplib_dict_headers(self):
+        _patch_httplib()
+        if sys.version_info >= (3, 0, 0):
+            from http.client import HTTPSConnection
+        else:
+            from httplib import HTTPSConnection
+
+        headers = MagicMock(spec=dict)
+        headers["fake-header"] = "fake-value"
+
+        conn = HTTPSConnection("www.datadoghq.com")
+        conn.request("GET", "/", headers=headers)
+        conn.getresponse()
+
+        self.mock_get_dd_trace_context.assert_called()
+        headers.update.assert_called()
+
+    def test_patch_httplib_dict_like_headers(self):
+        _patch_httplib()
+        if sys.version_info >= (3, 0, 0):
+            from http.client import HTTPSConnection
+            from collections.abc import MutableMapping
+        else:
+            from httplib import HTTPSConnection
+            from collections import MutableMapping
+
+        headers = MagicMock(spec=MutableMapping)
+        headers["fake-header"] = "fake-value"
+
+        conn = HTTPSConnection("www.datadoghq.com")
+        conn.request("GET", "/", headers=headers)
+        conn.getresponse()
+
         self.mock_get_dd_trace_context.assert_called()
+        headers.update.assert_called()
 
     def test_patch_requests(self):
         _ensure_patch_requests()