Add another test to test eventbridge extraction

datadog_lambda/tracing.py

@@ -239,7 +239,6 @@ def extract_dd_trace_context(event, lambda_context, extractor=None):
  """
  global dd_trace_context
  trace_context_source = None
-
  event_source = parse_event_source(event)
 
  if extractor is not None:
@@ -260,7 +259,7 @@ def extract_dd_trace_context(event, lambda_context, extractor=None):
  parent_id,
  sampling_priority,
  ) = extract_context_from_sqs_event_or_context(event, lambda_context)
- elif event_source == EventTypes.EVENTBRIDGE:
+ elif event_source.equals(EventTypes.EVENTBRIDGE):
  (
  trace_id,
  parent_id,

tests/event_samples/eventbridge-custom.json

@@ -10,6 +10,7 @@
  "detail": {
  "foo": "bar",
  "x-datadog-trace-id": "12345",
- "x-datadog-parent-id": "67890"
+ "x-datadog-parent-id": "67890",
+ "x-datadog-sampling-priority": "2"
  }
 }

tests/test_tracing.py

@@ -940,4 +940,14 @@ def test_extract_context_from_eventbridge_event(self):
  trace, parent, sampling = extract_context_from_eventbridge_event(event, ctx)
  self.assertEqual(trace, "12345")
  self.assertEqual(parent, "67890"),
- self.assertEqual(sampling, None)
+ self.assertEqual(sampling, "2")
+
+ def test_extract_dd_trace_context_for_eventbridge(self):
+ event_sample_source = "eventbridge-custom"
+ test_file = event_samples + event_sample_source + ".json"
+ with open(test_file, "r") as event:
+ event = json.load(event)
+ ctx = get_mock_context()
+ context, source = extract_dd_trace_context(event, ctx)
+ self.assertEqual(context["trace-id"], "12345")
+ self.assertEqual(context["parent-id"], "67890")