forked from DataDog/datadog-lambda-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_patch.py
More file actions
102 lines (81 loc) · 3.61 KB
/
test_patch.py
File metadata and controls
102 lines (81 loc) · 3.61 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
import sys
import unittest
try:
from unittest.mock import patch, MagicMock
except ImportError:
from mock import patch, MagicMock
from datadog_lambda.patch import _patch_httplib, _ensure_patch_requests
from datadog_lambda.constants import TraceHeader
class TestPatchHTTPClients(unittest.TestCase):
def setUp(self):
patcher = patch("datadog_lambda.patch.get_dd_trace_context")
self.mock_get_dd_trace_context = patcher.start()
self.mock_get_dd_trace_context.return_value = {
TraceHeader.TRACE_ID: "123",
TraceHeader.PARENT_ID: "321",
TraceHeader.SAMPLING_PRIORITY: "2",
}
self.addCleanup(patcher.stop)
def test_patch_httplib(self):
_patch_httplib()
if sys.version_info >= (3, 0, 0):
from http.client import HTTPSConnection
else:
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()
import requests
r = requests.get("https://www.datadoghq.com/")
self.mock_get_dd_trace_context.assert_called()
self.assertEqual(r.request.headers[TraceHeader.TRACE_ID], "123")
self.assertEqual(r.request.headers[TraceHeader.PARENT_ID], "321")
self.assertEqual(r.request.headers[TraceHeader.SAMPLING_PRIORITY], "2")
def test_patch_requests_with_headers(self):
_ensure_patch_requests()
import requests
r = requests.get("https://www.datadoghq.com/", headers={"key": "value"})
self.mock_get_dd_trace_context.assert_called()
self.assertEqual(r.request.headers["key"], "value")
self.assertEqual(r.request.headers[TraceHeader.TRACE_ID], "123")
self.assertEqual(r.request.headers[TraceHeader.PARENT_ID], "321")
self.assertEqual(r.request.headers[TraceHeader.SAMPLING_PRIORITY], "2")
def test_patch_requests_with_headers_none(self):
_ensure_patch_requests()
import requests
r = requests.get("https://www.datadoghq.com/", headers=None)
self.mock_get_dd_trace_context.assert_called()
self.assertEqual(r.request.headers[TraceHeader.TRACE_ID], "123")
self.assertEqual(r.request.headers[TraceHeader.PARENT_ID], "321")
self.assertEqual(r.request.headers[TraceHeader.SAMPLING_PRIORITY], "2")