Skip to content

Commit ffd0be5

Browse files
purple4reinapurple4reina
committed
Adds testing for w3c traceparent extraction.
1 parent 1c97b37 commit ffd0be5

File tree

1 file changed

+234
-2
lines changed

1 file changed

+234
-2
lines changed

‎tests/test_tracing.py

Lines changed: 234 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import unittest
2+
import functools
23
import json
34
import os
45
import copy
@@ -29,7 +30,6 @@
2930
_convert_xray_entity_id,
3031
_convert_xray_sampling,
3132
InferredSpanInfo,
32-
extract_context_from_eventbridge_event,
3333
create_service_mapping,
3434
determine_service_name,
3535
service_mapping as global_service_mapping,
@@ -96,6 +96,27 @@ def get_mock_context(
9696
return lambda_context
9797

9898

99+
def with_trace_propagation_style(style):
100+
style_list = list(style.split(","))
101+
102+
def _wrapper(fn):
103+
@functools.wraps(fn)
104+
def _wrap(*args, **kwargs):
105+
return fn(*args, **kwargs)
106+
from ddtrace.propagation.http import config
107+
108+
orig_style = config._propagation_style_extract
109+
config._propagation_style_extract = style_list
110+
try:
111+
return fn(*args, **kwargs)
112+
finally:
113+
config._propagation_style_extract = orig_style
114+
115+
return _wrap
116+
117+
return _wrapper
118+
119+
99120
class TestExtractAndGetDDTraceContext(unittest.TestCase):
100121
def setUp(self):
101122
global dd_tracing_enabled
@@ -161,6 +182,7 @@ def test_with_incomplete_datadog_trace_headers(self):
161182
self.assertEqual(ctx, expected_context)
162183
self.assertEqual(get_dd_trace_context(), expected_context)
163184

185+
@with_trace_propagation_style("datadog")
164186
def test_with_complete_datadog_trace_headers(self):
165187
lambda_ctx = get_mock_context()
166188
ctx, source, event_source = extract_dd_trace_context(
@@ -191,6 +213,45 @@ def test_with_complete_datadog_trace_headers(self):
191213
expected_context,
192214
)
193215

216+
@with_trace_propagation_style("tracecontext")
217+
def test_with_w3c_trace_headers(self):
218+
lambda_ctx = get_mock_context()
219+
ctx, source, event_source = extract_dd_trace_context(
220+
{
221+
"headers": {
222+
"traceparent": "00-0000000000000000000000000000007b-0000000000000141-01",
223+
"tracestate": "dd=s:2;t.dm:-0",
224+
}
225+
},
226+
lambda_ctx,
227+
)
228+
self.assertEqual(source, "event")
229+
expected_context = Context(
230+
trace_id=123,
231+
span_id=321,
232+
sampling_priority=2,
233+
meta={
234+
"traceparent": "00-0000000000000000000000000000007b-0000000000000141-01",
235+
"tracestate": "dd=s:2;t.dm:-0",
236+
"_dd.p.dm": "-0",
237+
},
238+
)
239+
self.assertEqual(ctx, expected_context)
240+
self.assertEqual(
241+
get_dd_trace_context(),
242+
Context(
243+
trace_id=123,
244+
span_id=fake_xray_header_value_parent_decimal,
245+
sampling_priority=2,
246+
),
247+
)
248+
create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY)
249+
self.mock_send_segment.assert_called()
250+
self.mock_send_segment.assert_called_with(
251+
XraySubsegment.TRACE_KEY,
252+
expected_context,
253+
)
254+
194255
def test_with_extractor_function(self):
195256
def extractor_foo(event, context):
196257
foo = event.get("foo", {})
@@ -256,6 +317,7 @@ def extractor_raiser(event, context):
256317
self.assertEqual(ctx, expected_context)
257318
self.assertEqual(get_dd_trace_context(), expected_context)
258319

320+
@with_trace_propagation_style("datadog")
259321
def test_with_sqs_distributed_datadog_trace_data(self):
260322
lambda_ctx = get_mock_context()
261323
sqs_event = {
@@ -311,6 +373,67 @@ def test_with_sqs_distributed_datadog_trace_data(self):
311373
expected_context,
312374
)
313375

376+
@with_trace_propagation_style("tracecontext")
377+
def test_with_sqs_distributed_w3c_trace_data(self):
378+
lambda_ctx = get_mock_context()
379+
sqs_event = {
380+
"Records": [
381+
{
382+
"messageId": "059f36b4-87a3-44ab-83d2-661975830a7d",
383+
"receiptHandle": "AQEBwJnKyrHigUMZj6rYigCgxlaS3SLy0a...",
384+
"body": "Test message.",
385+
"attributes": {
386+
"ApproximateReceiveCount": "1",
387+
"SentTimestamp": "1545082649183",
388+
"SenderId": "AIDAIENQZJOLO23YVJ4VO",
389+
"ApproximateFirstReceiveTimestamp": "1545082649185",
390+
},
391+
"messageAttributes": {
392+
"_datadog": {
393+
"stringValue": json.dumps(
394+
{
395+
"traceparent": "00-0000000000000000000000000000007b-0000000000000141-01",
396+
"tracestate": "dd=s:2;t.dm:-0",
397+
}
398+
),
399+
"dataType": "String",
400+
}
401+
},
402+
"md5OfBody": "e4e68fb7bd0e697a0ae8f1bb342846b3",
403+
"eventSource": "aws:sqs",
404+
"eventSourceARN": "arn:aws:sqs:us-east-2:123456789012:my-queue",
405+
"awsRegion": "us-east-2",
406+
}
407+
]
408+
}
409+
ctx, source, event_source = extract_dd_trace_context(sqs_event, lambda_ctx)
410+
self.assertEqual(source, "event")
411+
expected_context = Context(
412+
trace_id=123,
413+
span_id=321,
414+
sampling_priority=2,
415+
meta={
416+
"traceparent": "00-0000000000000000000000000000007b-0000000000000141-01",
417+
"tracestate": "dd=s:2;t.dm:-0",
418+
"_dd.p.dm": "-0",
419+
},
420+
)
421+
self.assertEqual(ctx, expected_context)
422+
self.assertEqual(
423+
get_dd_trace_context(),
424+
Context(
425+
trace_id=123,
426+
span_id=fake_xray_header_value_parent_decimal,
427+
sampling_priority=2,
428+
),
429+
)
430+
create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY)
431+
self.mock_send_segment.assert_called_with(
432+
XraySubsegment.TRACE_KEY,
433+
expected_context,
434+
)
435+
436+
@with_trace_propagation_style("datadog")
314437
def test_with_legacy_client_context_datadog_trace_data(self):
315438
lambda_ctx = get_mock_context(
316439
custom={
@@ -344,6 +467,45 @@ def test_with_legacy_client_context_datadog_trace_data(self):
344467
expected_context,
345468
)
346469

470+
@with_trace_propagation_style("tracecontext")
471+
def test_with_legacy_client_context_w3c_trace_data(self):
472+
lambda_ctx = get_mock_context(
473+
custom={
474+
"_datadog": {
475+
"traceparent": "00-0000000000000000000000000000029a-0000000000000309-01",
476+
"tracestate": "dd=s:1;t.dm:-0",
477+
}
478+
}
479+
)
480+
ctx, source, event_source = extract_dd_trace_context({}, lambda_ctx)
481+
self.assertEqual(source, "event")
482+
expected_context = Context(
483+
trace_id=666,
484+
span_id=777,
485+
sampling_priority=1,
486+
meta={
487+
"traceparent": "00-0000000000000000000000000000029a-0000000000000309-01",
488+
"tracestate": "dd=s:1;t.dm:-0",
489+
"_dd.p.dm": "-0",
490+
},
491+
)
492+
self.assertEqual(ctx, expected_context)
493+
self.assertEqual(
494+
get_dd_trace_context(),
495+
Context(
496+
trace_id=666,
497+
span_id=fake_xray_header_value_parent_decimal,
498+
sampling_priority=1,
499+
),
500+
)
501+
create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY)
502+
self.mock_send_segment.assert_called()
503+
self.mock_send_segment.assert_called_with(
504+
XraySubsegment.TRACE_KEY,
505+
expected_context,
506+
)
507+
508+
@with_trace_propagation_style("datadog")
347509
def test_with_new_client_context_datadog_trace_data(self):
348510
lambda_ctx = get_mock_context(
349511
custom={
@@ -375,6 +537,42 @@ def test_with_new_client_context_datadog_trace_data(self):
375537
expected_context,
376538
)
377539

540+
@with_trace_propagation_style("tracecontext")
541+
def test_with_new_client_context_w3c_trace_data(self):
542+
lambda_ctx = get_mock_context(
543+
custom={
544+
"traceparent": "00-0000000000000000000000000000029a-0000000000000309-01",
545+
"tracestate": "dd=s:1;t.dm:-0",
546+
}
547+
)
548+
ctx, source, event_source = extract_dd_trace_context({}, lambda_ctx)
549+
self.assertEqual(source, "event")
550+
expected_context = Context(
551+
trace_id=666,
552+
span_id=777,
553+
sampling_priority=1,
554+
meta={
555+
"traceparent": "00-0000000000000000000000000000029a-0000000000000309-01",
556+
"tracestate": "dd=s:1;t.dm:-0",
557+
"_dd.p.dm": "-0",
558+
},
559+
)
560+
self.assertEqual(ctx, expected_context)
561+
self.assertEqual(
562+
get_dd_trace_context(),
563+
Context(
564+
trace_id=666,
565+
span_id=fake_xray_header_value_parent_decimal,
566+
sampling_priority=1,
567+
),
568+
)
569+
create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY)
570+
self.mock_send_segment.assert_called()
571+
self.mock_send_segment.assert_called_with(
572+
XraySubsegment.TRACE_KEY,
573+
expected_context,
574+
)
575+
378576
def test_with_complete_datadog_trace_headers_with_mixed_casing(self):
379577
lambda_ctx = get_mock_context()
380578
extract_dd_trace_context(
@@ -417,6 +615,40 @@ def test_with_complete_datadog_trace_headers_with_trigger_tags(self):
417615
]
418616
)
419617

618+
def test_step_function_trace_data(self):
619+
lambda_ctx = get_mock_context()
620+
sqs_event = {
621+
"Execution": {
622+
"Id": "665c417c-1237-4742-aaca-8b3becbb9e75",
623+
},
624+
"StateMachine": {},
625+
"State": {
626+
"Name": "my-awesome-state",
627+
"EnteredTime": "Mon Nov 13 12:43:33 PST 2023",
628+
},
629+
}
630+
ctx, source, event_source = extract_dd_trace_context(sqs_event, lambda_ctx)
631+
self.assertEqual(source, "event")
632+
expected_context = Context(
633+
trace_id="1074655265866231755",
634+
span_id="4776286484851030060",
635+
sampling_priority=1,
636+
)
637+
self.assertEqual(ctx, expected_context)
638+
self.assertEqual(
639+
get_dd_trace_context(),
640+
Context(
641+
trace_id="1074655265866231755",
642+
span_id=fake_xray_header_value_parent_decimal,
643+
sampling_priority=1,
644+
),
645+
)
646+
create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY)
647+
self.mock_send_segment.assert_called_with(
648+
XraySubsegment.TRACE_KEY,
649+
expected_context,
650+
)
651+
420652

421653
class TestXRayContextConversion(unittest.TestCase):
422654
def test_convert_xray_trace_id(self):
@@ -1793,7 +2025,7 @@ def test_extract_context_from_eventbridge_event(self):
17932025
with open(test_file, "r") as event:
17942026
event = json.load(event)
17952027
ctx = get_mock_context()
1796-
context = extract_context_from_eventbridge_event(event, ctx)
2028+
context, source, event_type = extract_dd_trace_context(event, ctx)
17972029
self.assertEqual(context.trace_id, 12345)
17982030
self.assertEqual(context.span_id, 67890),
17992031
self.assertEqual(context.sampling_priority, 2)

0 commit comments

Comments
 (0)