Ensure dummy metadata segment is json serializable.

Fixes error `TypeError: Object of type Context is not JSON serializable`
#402

Author: purple4reina

datadog_lambda/wrapper.py

@@ -281,7 +281,12 @@ def _before(self, event, context):
             # Create a Datadog X-Ray subsegment with the trace context
             if dd_context and trace_context_source == TraceContextSource.EVENT:
                 create_dd_dummy_metadata_subsegment(
-                    dd_context, XraySubsegment.TRACE_KEY
+                    {
+                        'trace_id': dd_context.trace_id,
+                        'parent_id': dd_context.span_id,
+                        'sampling_priority': dd_context.sampling_priority,
+                    },
+                    XraySubsegment.TRACE_KEY,
                 )
 
             if dd_tracing_enabled:

tests/test_wrapper.py

@@ -9,7 +9,7 @@
 import datadog_lambda.wrapper as wrapper
 from datadog_lambda.metric import lambda_metric
 from datadog_lambda.thread_stats_writer import ThreadStatsWriter
-from ddtrace import Span
+from ddtrace import Span, tracer
 
 
 def get_mock_context(
@@ -580,3 +580,38 @@ def test_some_envs_should_depend_on_dd_tracing_enabled(self):
         self.assertFalse(decorator.make_inferred_span)
         self.assertFalse(decorator.encode_authorizer_context)
         self.assertFalse(decorator.decode_authorizer_context)
+
+class TestLambdaWrapperWithTraceContext(unittest.TestCase):
+    xray_daemon_envvar = "localhost:1234"
+    xray_trace_envvar = (
+        "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1;Lineage=c6c5b1b9:0"
+    )
+
+    @patch("os.environ", {
+        "AWS_XRAY_DAEMON_ADDRESS": xray_daemon_envvar,
+        "_X_AMZN_TRACE_ID": xray_trace_envvar,
+    })
+    def test_event_bridge_sqs_payload(self):
+        wrapper.dd_tracing_enabled = True
+
+        def handler(event, context):
+            return tracer.current_trace_context()
+
+        wrapped_handler = wrapper.datadog_lambda_wrapper(handler)
+
+        trace_id = "2a4abb4251b6ec1d"
+        event = {
+            "headers": {
+                "traceparent": f"00-0000000000000000{trace_id}-74d86a8e4b55daf4-01",
+                "tracestate": "dd=s:1;t.dm:-1",
+            },
+        }
+        context = get_mock_context()
+
+        result = wrapped_handler(event, context)
+        aws_lambda_span = wrapped_handler.span
+
+        self.assertIsNotNone(result)
+        self.assertEqual(result.trace_id, int(trace_id, 16))
+        self.assertEqual(result.span_id, aws_lambda_span.span_id)
+        self.assertEqual(result.sampling_priority, 1)