Refactored code, added more unit tests

Author: rithikanarayan

datadog_lambda/tracing.py

@@ -209,6 +209,20 @@ def extract_context_from_http_event_or_context(
     return context
 
 
+def extract_context_from_request_header_or_context(event, lambda_context, event_source):
+    request = event.get("request")
+    if isinstance(request, (set, dict)) and "headers" in request:
+        context = extract_context_from_http_event_or_context(
+            request,
+            lambda_context,
+            event_source,
+            decode_authorizer_context=False,
+        )
+    else:
+        context = extract_context_from_lambda_context(lambda_context)
+    return context
+
+
 def create_sns_event(message):
     return {
         "Records": [
@@ -628,16 +642,9 @@ def extract_dd_trace_context(
     if extractor is not None:
         context = extract_context_custom_extractor(extractor, event, lambda_context)
     elif isinstance(event, (set, dict)) and "request" in event:
-        request = event.get("request")
-        if isinstance(request, (set, dict)) and "headers" in request:
-            context = extract_context_from_http_event_or_context(
-                request,
-                lambda_context,
-                event_source,
-                decode_authorizer_context=False,
-            )
-        else:
-            context = extract_context_from_lambda_context(lambda_context)
+        context = extract_context_from_request_header_or_context(
+            event, lambda_context, event_source
+        )
     elif isinstance(event, (set, dict)) and "headers" in event:
         context = extract_context_from_http_event_or_context(
             event, lambda_context, event_source, decode_authorizer_context

env/pyvenv.cfg

@@ -0,0 +1,3 @@
+home = /Users/rithika.narayan/.pyenv/versions/3.10.13/bin
+include-system-site-packages = false
+version = 3.10.13

tests/event_samples/rum-appsync-no-headers.json

@@ -0,0 +1,14 @@
+{
+    "identity": "None",
+    "info": {
+        "fieldName": "getItems",
+        "parentTypeName": "Query",
+        "selectionSetGraphQL": "{\n  id\n}",
+        "selectionSetList":["id"]
+    }, 
+    "prev": "None",
+    "request": { 
+        "domainName": "None"
+    },
+    "source": "None"
+}

tests/event_samples/rum-appsync-request-not-dict.json

@@ -0,0 +1,12 @@
+{
+    "identity": "None",
+    "info": {
+        "fieldName": "getItems",
+        "parentTypeName": "Query",
+        "selectionSetGraphQL": "{\n  id\n}",
+        "selectionSetList":["id"]
+    }, 
+    "prev": "None",
+    "request": "hello",
+    "source": "None"
+}

tests/test_tracing.py

@@ -187,6 +187,8 @@ def _wrap(*args, **kwargs):
             sampling_priority=1,
         ),
     ),
+    ("rum-appsync-no-headers", None),
+    ("rum-appsync-request-not-dict", None), 
     ("s3", None),
     (
         "sns-b64-msg-attribute",