@@ -272,7 +272,7 @@ def test_with_sqs_distributed_datadog_trace_data(self):
272
272
XraySubsegment .NAMESPACE ,
273
273
)
274
274
275
- def test_with_client_context_datadog_trace_data (self ):
275
+ def test_with_legacy_client_context_datadog_trace_data (self ):
276
276
lambda_ctx = get_mock_context (
277
277
custom = {
278
278
"_datadog" : {
@@ -309,6 +309,41 @@ def test_with_client_context_datadog_trace_data(self):
309
309
XraySubsegment .NAMESPACE ,
310
310
)
311
311
312
+ def test_with_new_client_context_datadog_trace_data (self ):
313
+ lambda_ctx = get_mock_context (
314
+ custom = {
315
+ TraceHeader .TRACE_ID : "666" ,
316
+ TraceHeader .PARENT_ID : "777" ,
317
+ TraceHeader .SAMPLING_PRIORITY : "1" ,
318
+ }
319
+ )
320
+ ctx , source = extract_dd_trace_context ({}, lambda_ctx )
321
+ self .assertEqual (source , "event" )
322
+ self .assertDictEqual (
323
+ ctx ,
324
+ {
325
+ "trace-id" : "666" ,
326
+ "parent-id" : "777" ,
327
+ "sampling-priority" : "1" ,
328
+ },
329
+ )
330
+ self .assertDictEqual (
331
+ get_dd_trace_context (),
332
+ {
333
+ TraceHeader .TRACE_ID : "666" ,
334
+ TraceHeader .PARENT_ID : "65535" ,
335
+ TraceHeader .SAMPLING_PRIORITY : "1" ,
336
+ },
337
+ )
338
+ create_dd_dummy_metadata_subsegment (ctx , XraySubsegment .TRACE_KEY )
339
+ self .mock_xray_recorder .begin_subsegment .assert_called ()
340
+ self .mock_xray_recorder .end_subsegment .assert_called ()
341
+ self .mock_current_subsegment .put_metadata .assert_called_with (
342
+ XraySubsegment .TRACE_KEY ,
343
+ {"trace-id" : "666" , "parent-id" : "777" , "sampling-priority" : "1" },
344
+ XraySubsegment .NAMESPACE ,
345
+ )
346
+
312
347
def test_with_complete_datadog_trace_headers_with_mixed_casing (self ):
313
348
lambda_ctx = get_mock_context ()
314
349
extract_dd_trace_context (
0 commit comments