fix the str response case

Author: joeyzhao2018

datadog_lambda/tracing.py

@@ -344,6 +344,11 @@ def extract_context_custom_extractor(extractor, event, lambda_context):
 
         return None, None, None
 
+def is_authorizer_response(response) -> bool:
+    try:
+        return response["principalId"] and response["policyDocument"]
+    except:
+        return False
 
 def get_injected_authorizer_data(event, is_http_api) -> dict:
     try:

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,
@@ -275,12 +276,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,19 @@ 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")
+
+    def test_different_return_type_no_error(self):
+        TEST_RESULTS = ["a str to return", 42, {'value': 42}, ['A', 42]]
+        import traceback as _importedTraceback
+        _importedTraceback.print_exc = MagicMock()
+        mock_context = get_mock_context()
+        for test_result in TEST_RESULTS:
+
+            @datadog_lambda_wrapper
+            def return_a_str(event, context):
+                return test_result
+
+            result = return_a_str({}, mock_context)
+            self.assertEquals(result, test_result)
+            _importedTraceback.print_exc.assert_not_called()
+