Skip to content

Commit 4669117

Browse files
purple4reinapurple4reina
authored
Parametrized trace context extraction tests. (#456)
* Parametrized trace context extraction tests. * Linting.
1 parent 92990ae commit 4669117

File tree

1 file changed

+161
-129
lines changed

1 file changed

+161
-129
lines changed

‎tests/test_tracing.py‎

Lines changed: 161 additions & 129 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,167 @@ def _wrap(*args, **kwargs):
9696
return _wrapper
9797

9898

99+
_test_extract_dd_trace_context = (
100+
("api-gateway", Context(trace_id=12345, span_id=67890, sampling_priority=2)),
101+
(
102+
"api-gateway-no-apiid",
103+
Context(trace_id=12345, span_id=67890, sampling_priority=2),
104+
),
105+
(
106+
"api-gateway-non-proxy",
107+
Context(trace_id=12345, span_id=67890, sampling_priority=2),
108+
),
109+
(
110+
"api-gateway-non-proxy-async",
111+
Context(trace_id=12345, span_id=67890, sampling_priority=2),
112+
),
113+
(
114+
"api-gateway-websocket-connect",
115+
Context(trace_id=12345, span_id=67890, sampling_priority=2),
116+
),
117+
(
118+
"api-gateway-websocket-default",
119+
Context(trace_id=12345, span_id=67890, sampling_priority=2),
120+
),
121+
(
122+
"api-gateway-websocket-disconnect",
123+
Context(trace_id=12345, span_id=67890, sampling_priority=2),
124+
),
125+
(
126+
"authorizer-request-api-gateway-v1",
127+
Context(
128+
trace_id=13478705995797221209,
129+
span_id=8471288263384216896,
130+
sampling_priority=1,
131+
),
132+
),
133+
("authorizer-request-api-gateway-v1-cached", None),
134+
(
135+
"authorizer-request-api-gateway-v2",
136+
Context(
137+
trace_id=14356983619852933354,
138+
span_id=12658621083505413809,
139+
sampling_priority=1,
140+
),
141+
),
142+
("authorizer-request-api-gateway-v2-cached", None),
143+
(
144+
"authorizer-request-api-gateway-websocket-connect",
145+
Context(
146+
trace_id=5351047404834723189,
147+
span_id=18230460631156161837,
148+
sampling_priority=1,
149+
),
150+
),
151+
("authorizer-request-api-gateway-websocket-message", None),
152+
(
153+
"authorizer-token-api-gateway-v1",
154+
Context(
155+
trace_id=17874798268144902712,
156+
span_id=16184667399315372101,
157+
sampling_priority=1,
158+
),
159+
),
160+
("authorizer-token-api-gateway-v1-cached", None),
161+
("cloudfront", None),
162+
("cloudwatch-events", None),
163+
("cloudwatch-logs", None),
164+
("custom", None),
165+
("dynamodb", None),
166+
("eventbridge-custom", Context(trace_id=12345, span_id=67890, sampling_priority=2)),
167+
(
168+
"eventbridge-sqs",
169+
Context(
170+
trace_id=7379586022458917877,
171+
span_id=2644033662113726488,
172+
sampling_priority=1,
173+
),
174+
),
175+
("http-api", Context(trace_id=12345, span_id=67890, sampling_priority=2)),
176+
(
177+
"kinesis",
178+
Context(
179+
trace_id=4948377316357291421,
180+
span_id=2876253380018681026,
181+
sampling_priority=1,
182+
),
183+
),
184+
(
185+
"kinesis-batch",
186+
Context(
187+
trace_id=4948377316357291421,
188+
span_id=2876253380018681026,
189+
sampling_priority=1,
190+
),
191+
),
192+
("lambda-url", None),
193+
("s3", None),
194+
(
195+
"sns-b64-msg-attribute",
196+
Context(
197+
trace_id=4948377316357291421,
198+
span_id=6746998015037429512,
199+
sampling_priority=1,
200+
),
201+
),
202+
(
203+
"sns-batch",
204+
Context(
205+
trace_id=4948377316357291421,
206+
span_id=6746998015037429512,
207+
sampling_priority=1,
208+
),
209+
),
210+
(
211+
"sns-string-msg-attribute",
212+
Context(
213+
trace_id=4948377316357291421,
214+
span_id=6746998015037429512,
215+
sampling_priority=1,
216+
),
217+
),
218+
(
219+
"sqs-batch",
220+
Context(
221+
trace_id=2684756524522091840,
222+
span_id=7431398482019833808,
223+
sampling_priority=1,
224+
),
225+
),
226+
(
227+
"sqs-java-upstream",
228+
Context(
229+
trace_id=7925498337868555493,
230+
span_id=5245570649555658903,
231+
sampling_priority=1,
232+
),
233+
),
234+
(
235+
"sqs-string-msg-attribute",
236+
Context(
237+
trace_id=2684756524522091840,
238+
span_id=7431398482019833808,
239+
sampling_priority=1,
240+
),
241+
),
242+
({"headers": None}, None),
243+
)
244+
245+
246+
@pytest.mark.parametrize("event,expect", _test_extract_dd_trace_context)
247+
def test_extract_dd_trace_context(event, expect):
248+
if isinstance(event, str):
249+
with open(f"{event_samples}{event}.json") as f:
250+
event = json.load(f)
251+
ctx = get_mock_context()
252+
253+
actual, _, _ = extract_dd_trace_context(event, ctx)
254+
assert (expect is None) is (actual is None)
255+
assert (expect is None) or actual.trace_id == expect.trace_id
256+
assert (expect is None) or actual.span_id == expect.span_id
257+
assert (expect is None) or actual.sampling_priority == expect.sampling_priority
258+
259+
99260
class TestExtractAndGetDDTraceContext(unittest.TestCase):
100261
def setUp(self):
101262
global dd_tracing_enabled
@@ -1773,127 +1934,6 @@ def test_create_inferred_span(mock_span_finish, source, expect):
17731934

17741935

17751936
class TestInferredSpans(unittest.TestCase):
1776-
def test_extract_context_from_eventbridge_event(self):
1777-
event_sample_source = "eventbridge-custom"
1778-
test_file = event_samples + event_sample_source + ".json"
1779-
with open(test_file, "r") as event:
1780-
event = json.load(event)
1781-
ctx = get_mock_context()
1782-
context, source, event_type = extract_dd_trace_context(event, ctx)
1783-
self.assertEqual(context.trace_id, 12345)
1784-
self.assertEqual(context.span_id, 67890),
1785-
self.assertEqual(context.sampling_priority, 2)
1786-
1787-
def test_extract_dd_trace_context_for_eventbridge(self):
1788-
event_sample_source = "eventbridge-custom"
1789-
test_file = event_samples + event_sample_source + ".json"
1790-
with open(test_file, "r") as event:
1791-
event = json.load(event)
1792-
ctx = get_mock_context()
1793-
context, source, event_type = extract_dd_trace_context(event, ctx)
1794-
self.assertEqual(context.trace_id, 12345)
1795-
self.assertEqual(context.span_id, 67890)
1796-
1797-
def test_extract_context_from_eventbridge_sqs_event(self):
1798-
event_sample_source = "eventbridge-sqs"
1799-
test_file = event_samples + event_sample_source + ".json"
1800-
with open(test_file, "r") as event:
1801-
event = json.load(event)
1802-
1803-
ctx = get_mock_context()
1804-
context, source, event_type = extract_dd_trace_context(event, ctx)
1805-
self.assertEqual(context.trace_id, 7379586022458917877)
1806-
self.assertEqual(context.span_id, 2644033662113726488)
1807-
self.assertEqual(context.sampling_priority, 1)
1808-
1809-
def test_extract_context_from_sqs_event_with_string_msg_attr(self):
1810-
event_sample_source = "sqs-string-msg-attribute"
1811-
test_file = event_samples + event_sample_source + ".json"
1812-
with open(test_file, "r") as event:
1813-
event = json.load(event)
1814-
ctx = get_mock_context()
1815-
context, source, event_type = extract_dd_trace_context(event, ctx)
1816-
self.assertEqual(context.trace_id, 2684756524522091840)
1817-
self.assertEqual(context.span_id, 7431398482019833808)
1818-
self.assertEqual(context.sampling_priority, 1)
1819-
1820-
def test_extract_context_from_sqs_batch_event(self):
1821-
event_sample_source = "sqs-batch"
1822-
test_file = event_samples + event_sample_source + ".json"
1823-
with open(test_file, "r") as event:
1824-
event = json.load(event)
1825-
ctx = get_mock_context()
1826-
context, source, event_source = extract_dd_trace_context(event, ctx)
1827-
self.assertEqual(context.trace_id, 2684756524522091840)
1828-
self.assertEqual(context.span_id, 7431398482019833808)
1829-
self.assertEqual(context.sampling_priority, 1)
1830-
1831-
def test_extract_context_from_sqs_java_upstream_event(self):
1832-
event_sample_source = "sqs-java-upstream"
1833-
test_file = event_samples + event_sample_source + ".json"
1834-
with open(test_file, "r") as event:
1835-
event = json.load(event)
1836-
ctx = get_mock_context()
1837-
context, source, event_type = extract_dd_trace_context(event, ctx)
1838-
self.assertEqual(context.trace_id, 7925498337868555493)
1839-
self.assertEqual(context.span_id, 5245570649555658903)
1840-
self.assertEqual(context.sampling_priority, 1)
1841-
1842-
def test_extract_context_from_sns_event_with_string_msg_attr(self):
1843-
event_sample_source = "sns-string-msg-attribute"
1844-
test_file = event_samples + event_sample_source + ".json"
1845-
with open(test_file, "r") as event:
1846-
event = json.load(event)
1847-
ctx = get_mock_context()
1848-
context, source, event_source = extract_dd_trace_context(event, ctx)
1849-
self.assertEqual(context.trace_id, 4948377316357291421)
1850-
self.assertEqual(context.span_id, 6746998015037429512)
1851-
self.assertEqual(context.sampling_priority, 1)
1852-
1853-
def test_extract_context_from_sns_event_with_b64_msg_attr(self):
1854-
event_sample_source = "sns-b64-msg-attribute"
1855-
test_file = event_samples + event_sample_source + ".json"
1856-
with open(test_file, "r") as event:
1857-
event = json.load(event)
1858-
ctx = get_mock_context()
1859-
context, source, event_source = extract_dd_trace_context(event, ctx)
1860-
self.assertEqual(context.trace_id, 4948377316357291421)
1861-
self.assertEqual(context.span_id, 6746998015037429512)
1862-
self.assertEqual(context.sampling_priority, 1)
1863-
1864-
def test_extract_context_from_sns_batch_event(self):
1865-
event_sample_source = "sns-batch"
1866-
test_file = event_samples + event_sample_source + ".json"
1867-
with open(test_file, "r") as event:
1868-
event = json.load(event)
1869-
ctx = get_mock_context()
1870-
context, source, event_source = extract_dd_trace_context(event, ctx)
1871-
self.assertEqual(context.trace_id, 4948377316357291421)
1872-
self.assertEqual(context.span_id, 6746998015037429512)
1873-
self.assertEqual(context.sampling_priority, 1)
1874-
1875-
def test_extract_context_from_kinesis_event(self):
1876-
event_sample_source = "kinesis"
1877-
test_file = event_samples + event_sample_source + ".json"
1878-
with open(test_file, "r") as event:
1879-
event = json.load(event)
1880-
ctx = get_mock_context()
1881-
context, source, event_source = extract_dd_trace_context(event, ctx)
1882-
self.assertEqual(context.trace_id, 4948377316357291421)
1883-
self.assertEqual(context.span_id, 2876253380018681026)
1884-
self.assertEqual(context.sampling_priority, 1)
1885-
1886-
def test_extract_context_from_kinesis_batch_event(self):
1887-
event_sample_source = "kinesis-batch"
1888-
test_file = event_samples + event_sample_source + ".json"
1889-
with open(test_file, "r") as event:
1890-
event = json.load(event)
1891-
ctx = get_mock_context()
1892-
context, source, event_source = extract_dd_trace_context(event, ctx)
1893-
self.assertEqual(context.trace_id, 4948377316357291421)
1894-
self.assertEqual(context.span_id, 2876253380018681026)
1895-
self.assertEqual(context.sampling_priority, 1)
1896-
18971937
@patch("datadog_lambda.tracing.submit_errors_metric")
18981938
def test_mark_trace_as_error_for_5xx_responses_getting_400_response_code(
18991939
self, mock_submit_errors_metric
@@ -1915,14 +1955,6 @@ def test_mark_trace_as_error_for_5xx_responses_sends_error_metric_and_set_error_
19151955
mock_submit_errors_metric.assert_called_once()
19161956
self.assertEqual(1, mock_span.error)
19171957

1918-
def test_no_error_with_nonetype_headers(self):
1919-
lambda_ctx = get_mock_context()
1920-
ctx, source, event_type = extract_dd_trace_context(
1921-
{"headers": None},
1922-
lambda_ctx,
1923-
)
1924-
self.assertEqual(ctx, None)
1925-
19261958

19271959
class TestStepFunctionsTraceContext(unittest.TestCase):
19281960
def test_deterministic_m5_hash(self):

0 commit comments

Comments
 (0)