Skip to content

Commit 3312a30

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

File tree

1 file changed

+233
-2
lines changed

1 file changed

+233
-2
lines changed

‎tests/test_tracing.py

Lines changed: 233 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,26 @@ 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+
from ddtrace.propagation.http import config
106+
107+
orig_style = config._propagation_style_extract
108+
config._propagation_style_extract = style_list
109+
try:
110+
return fn(*args, **kwargs)
111+
finally:
112+
config._propagation_style_extract = orig_style
113+
114+
return _wrap
115+
116+
return _wrapper
117+
118+
99119
class TestExtractAndGetDDTraceContext(unittest.TestCase):
100120
def setUp(self):
101121
global dd_tracing_enabled
@@ -161,6 +181,7 @@ def test_with_incomplete_datadog_trace_headers(self):
161181
self.assertEqual(ctx, expected_context)
162182
self.assertEqual(get_dd_trace_context(), expected_context)
163183

184+
@with_trace_propagation_style("datadog")
164185
def test_with_complete_datadog_trace_headers(self):
165186
lambda_ctx = get_mock_context()
166187
ctx, source, event_source = extract_dd_trace_context(
@@ -191,6 +212,45 @@ def test_with_complete_datadog_trace_headers(self):
191212
expected_context,
192213
)
193214

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

319+
@with_trace_propagation_style("datadog")
259320
def test_with_sqs_distributed_datadog_trace_data(self):
260321
lambda_ctx = get_mock_context()
261322
sqs_event = {
@@ -311,6 +372,67 @@ def test_with_sqs_distributed_datadog_trace_data(self):
311372
expected_context,
312373
)
313374

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

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

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

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

421652
class TestXRayContextConversion(unittest.TestCase):
422653
def test_convert_xray_trace_id(self):
@@ -1793,7 +2024,7 @@ def test_extract_context_from_eventbridge_event(self):
17932024
with open(test_file, "r") as event:
17942025
event = json.load(event)
17952026
ctx = get_mock_context()
1796-
context = extract_context_from_eventbridge_event(event, ctx)
2027+
context, source, event_type = extract_dd_trace_context(event, ctx)
17972028
self.assertEqual(context.trace_id, 12345)
17982029
self.assertEqual(context.span_id, 67890),
17992030
self.assertEqual(context.sampling_priority, 2)

0 commit comments

Comments
 (0)