[BUGFIX] Fix how we look up the client context (#107)

Czechh · web-flow · commit a90a04549618 · 2021-01-07T14:14:31.000-05:00
* fix how we look up the client context

* fix lint
diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py
@@ -99,8 +99,8 @@ def extract_context_from_lambda_context(lambda_context):
     """
     client_context = lambda_context.client_context
 
-    if client_context and "custom" in client_context:
-        dd_data = client_context.get("custom", {}).get("_datadog", {})
+    if client_context and client_context.custom:
+        dd_data = client_context.custom.get("_datadog", {})
         trace_id = dd_data.get(TraceHeader.TRACE_ID)
         parent_id = dd_data.get(TraceHeader.PARENT_ID)
         sampling_priority = dd_data.get(TraceHeader.SAMPLING_PRIORITY)
diff --git a/tests/test_tracing.py b/tests/test_tracing.py
@@ -22,19 +22,24 @@
 function_arn = "arn:aws:lambda:us-west-1:123457598159:function:python-layer-test"
 
 
+class ClientContext(object):
+    def __init__(self, custom=None):
+        self.custom = custom
+
+
 def get_mock_context(
     aws_request_id="request-id-1",
     memory_limit_in_mb="256",
     invoked_function_arn=function_arn,
     function_version="1",
-    client_context={},
+    custom=None,
 ):
     lambda_context = MagicMock()
     lambda_context.aws_request_id = aws_request_id
     lambda_context.memory_limit_in_mb = memory_limit_in_mb
     lambda_context.invoked_function_arn = invoked_function_arn
     lambda_context.function_version = function_version
-    lambda_context.client_context = client_context
+    lambda_context.client_context = ClientContext(custom)
     return lambda_context
 
 
@@ -205,13 +210,11 @@ def test_with_sqs_distributed_datadog_trace_data(self):
 
     def test_with_client_context_datadog_trace_data(self):
         lambda_ctx = get_mock_context(
-            client_context={
-                "custom": {
-                    "_datadog": {
-                        TraceHeader.TRACE_ID: "666",
-                        TraceHeader.PARENT_ID: "777",
-                        TraceHeader.SAMPLING_PRIORITY: "1",
-                    }
+            custom={
+                "_datadog": {
+                    TraceHeader.TRACE_ID: "666",
+                    TraceHeader.PARENT_ID: "777",
+                    TraceHeader.SAMPLING_PRIORITY: "1",
                 }
             }
         )
