fix the str response case (#296)

Author: joeyzhao2018

datadog_lambda/tracing.py

@@ -345,6 +345,20 @@ def extract_context_custom_extractor(extractor, event, lambda_context):
         return None, None, None
 
 
+def is_authorizer_response(response) -> bool:
+    try:
+        return (
+            response is not None
+            and response["principalId"]
+            and response["policyDocument"]
+        )
+    except KeyError:
+        pass
+    except Exception as e:
+        logger.debug("unknown error while checking is_authorizer_response %s", e)
+    return False
+
+
 def get_injected_authorizer_data(event, is_http_api) -> dict:
     try:
         authorizer_headers = event.get("requestContext", {}).get("authorizer")

datadog_lambda/wrapper.py

@@ -37,6 +37,7 @@
     create_function_execution_span,
     create_inferred_span,
     InferredSpanInfo,
+    is_authorizer_response,
 )
 from datadog_lambda.trigger import (
     extract_trigger_tags,
@@ -280,12 +281,7 @@ def _after(self, event, context):
             if should_use_extension:
                 flush_extension()
 
-            if (
-                self.encode_authorizer_context
-                and self.response
-                and self.response.get("principalId")
-                and self.response.get("policyDocument")
-            ):
+            if self.encode_authorizer_context and is_authorizer_response(self.response):
                 self._inject_authorizer_span_headers(
                     event.get("requestContext", {}).get("requestId")
                 )

tests/test_wrapper.py

@@ -544,3 +544,17 @@ def lambda_handler(event, context):
         self.assertEquals(inject_data[TraceHeader.TRACE_ID], "456")
         self.assertEquals(inject_data[TraceHeader.SAMPLING_PRIORITY], "1")
         self.assertEquals(result["context"]["scope"], "still here")
+
+    @patch("traceback.print_exc")
+    def test_different_return_type_no_error(self, MockPrintExc):
+        TEST_RESULTS = ["a str to return", 42, {"value": 42}, ["A", 42], None]
+        mock_context = get_mock_context()
+        for test_result in TEST_RESULTS:
+
+            @datadog_lambda_wrapper
+            def return_type_test(event, context):
+                return test_result
+
+            result = return_type_test({}, mock_context)
+            self.assertEquals(result, test_result)
+            self.assertFalse(MockPrintExc.called)