added unit tests

Author: avedmala

datadog_lambda/tracing.py

@@ -371,7 +371,10 @@ def _deterministic_sha256_hash(s: str, part: str) -> int:
 
 
 def _parse_high_64_bits(trace_tags: str) -> str:
-    # todo: testme
+    """
+    Parse a list of trace tags such as [_dd.p.tid=66bcb5eb00000000,_dd.p.dm=-0] and return the value of the _dd.p.tid
+    tag or an empty string if not found.
+    """
     if trace_tags:
         for tag in trace_tags.split(","):
             if "_dd.p.tid=" in tag:
@@ -380,7 +383,6 @@ def _parse_high_64_bits(trace_tags: str) -> str:
     return ""
 
 
-# todo: testme
 def _generate_sfn_parent_id(context: dict) -> int:
     execution_id = context.get("Execution").get("Id")
     state_name = context.get("State").get("Name")
@@ -391,24 +393,31 @@ def _generate_sfn_parent_id(context: dict) -> int:
     )
 
 
-# take the higher 64 bits as _dd.p.tid tag and use hex to encode
-# [2:] to remove '0x' in the hex str
-# todo testme
-# returning 128 bits since 128bit traceId will be break up into
-# traditional traceId and _dd.p.tid tag
-# https://github.com/DataDog/dd-trace-py/blob/3e34d21cb9b5e1916e549047158cb119317b96ab/ddtrace/propagation/http.py#L232-L240
-def _generate_sfn_trace_id(pre_hash: str, part: str):
+def _generate_sfn_trace_id(execution_id: str, part: str):
+    """
+    Take the SHA-256 hash of the execution_id to calculate the trace ID. If the high 64 bits are specified, we take
+    those bits and use hex to encode it. We also remove the first two characters as they will be '0x in the hex string.
+
+    We care about full 128 bits because they will break up into traditional traceID and _dd.p.tid tag.
+    https://github.com/DataDog/dd-trace-py/blob/3e34d21cb9b5e1916e549047158cb119317b96ab/ddtrace/propagation/http.py#L232-L240
+    """
     if part == HIGHER_64_BITS:
-        return hex(_deterministic_sha256_hash(pre_hash, part))[2:]
-    return _deterministic_sha256_hash(pre_hash, part)
+        return hex(_deterministic_sha256_hash(execution_id, part))[2:]
+    return _deterministic_sha256_hash(execution_id, part)
 
 
 def extract_context_from_step_functions(event, lambda_context):
     """
     Only extract datadog trace context when Step Functions Context Object is injected
     into lambda's event dict.
+
+    If '_datadog' header is present, we have two cases:
+      1. Root is a Lambda and we use its traceID
+      2. Root is a SFN, and we use its executionARN to calculate the traceID
+    We calculate the parentID the same in both cases by using the parent SFN's context object.
+
+    Otherwise, we're dealing with the legacy case where we only have the parent SFN's context object.
     """
-    # todo: update docstring
     try:
         meta = {}
         dd_data = event.get("_datadog")

datadog_lambda/trigger.py

@@ -256,7 +256,6 @@ def parse_event_source_arn(source: _EventSource, event: dict, context: Any) -> s
     if source.event_type == EventTypes.CLOUDWATCH_EVENTS and event.get("resources"):
         return event.get("resources")[0]
 
-    # todo: testme
     # e.g. arn:aws:states:us-east-1:123456789012:stateMachine:stateMachineName
     if source.event_type == EventTypes.STEPFUNCTIONS:
         context = event

tests/event_samples/stepfunctions.json

@@ -0,0 +1,26 @@
+{
+  "Records": [
+    {
+      "_datadog": {
+        "Execution": {
+          "Id": "arn:aws:states:ca-central-1:425362996713:execution:MyStateMachine-wsx8chv4d:1356a963-42a5-48b0-ba3f-73bde559a50c",
+          "StartTime": "2024-11-13T16:46:47.715Z",
+          "Name": "1356a963-42a5-48b0-ba3f-73bde559a50c",
+          "RoleArn": "arn:aws:iam::425362996713:role/service-role/StepFunctions-MyStateMachine-wsx8chv4d-role-1su0fkfd3",
+          "RedriveCount": 0
+        },
+        "StateMachine": {
+          "Id": "arn:aws:states:ca-central-1:425362996713:stateMachine:MyStateMachine-wsx8chv4d",
+          "Name": "MyStateMachine-wsx8chv4d"
+        },
+        "State": {
+          "Name": "Lambda Invoke",
+          "EnteredTime": "2024-11-13T16:46:47.740Z",
+          "RetryCount": 0
+        },
+        "RootExecutionId": "arn:aws:states:ca-central-1:425362996713:execution:MyStateMachine-wsx8chv4d:1356a963-42a5-48b0-ba3f-73bde559a50c",
+        "serverless-version": "v2"
+      }
+    }
+  ]
+}

tests/test_trigger.py

@@ -230,6 +230,20 @@ def test_event_source_sqs(self):
             "arn:aws:sqs:eu-west-1:601427279990:InferredSpansQueueNode",
         )
 
+    def test_event_source_stepfunctions(self):
+        event_sample_source = "stepfunctions"
+        test_file = event_samples + event_sample_source + ".json"
+        with open(test_file, "r") as event:
+            event = json.load(event)
+        ctx = get_mock_context()
+        event_source = parse_event_source(event)
+        event_source_arn = get_event_source_arn(event_source, event, ctx)
+        self.assertEqual(event_source.to_string(), event_sample_source)
+        self.assertEqual(
+            event_source_arn,
+            "arn:aws:states:ca-central-1:425362996713:stateMachine:MyStateMachine-wsx8chv4d",
+        )
+
     def test_event_source_unsupported(self):
         event_sample_source = "custom"
         test_file = event_samples + event_sample_source + ".json"
@@ -485,6 +499,21 @@ def test_extract_trigger_tags_sqs(self):
             },
         )
 
+    def test_extract_trigger_tags_stepfunctions(self):
+        event_sample_source = "stepfunctions"
+        test_file = event_samples + event_sample_source + ".json"
+        ctx = get_mock_context()
+        with open(test_file, "r") as event:
+            event = json.load(event)
+        tags = extract_trigger_tags(event, ctx)
+        self.assertEqual(
+            tags,
+            {
+                "function_trigger.event_source": "stepfunctions",
+                "function_trigger.event_source_arn": "arn:aws:states:ca-central-1:425362996713:stateMachine:MyStateMachine-wsx8chv4d",
+            },
+        )
+
     def test_extract_trigger_tags_unsupported(self):
         event_sample_source = "custom"
         test_file = event_samples + event_sample_source + ".json"