Fix get status_code (#117)

datadog_lambda/trigger.py

@@ -207,7 +207,9 @@ def extract_http_status_code_tag(trigger_tags, response):
     if response is None:
         # Return a 502 status if no response is found
         status_code = "502"
-    elif response.get("statusCode"):
+    elif hasattr(response, "get"):
         status_code = response.get("statusCode")
+    elif hasattr(response, "status_code"):
+        status_code = response.status_code
 
     return status_code

tests/test_trigger.py

@@ -11,6 +11,7 @@
     parse_event_source,
     get_event_source_arn,
     extract_trigger_tags,
+    extract_http_status_code_tag,
 )
 
 event_samples = "tests/event_samples/"
@@ -328,3 +329,18 @@ def test_extract_trigger_tags_unsupported(self):
             event = json.load(event)
         tags = extract_trigger_tags(event, ctx)
         self.assertEqual(tags, {})
+
+
+class ExtractHTTPStatusCodeTag(unittest.TestCase):
+    def test_extract_http_status_code_tag_from_response_dict(self):
+        trigger_tags = {"function_trigger.event_source": "api-gateway"}
+        response = {"statusCode": 403}
+        status_code = extract_http_status_code_tag(trigger_tags, response)
+        self.assertEqual(status_code, 403)
+
+    def test_extract_http_status_code_tag_from_response_object(self):
+        trigger_tags = {"function_trigger.event_source": "api-gateway"}
+        response = MagicMock(spec=["status_code"])
+        response.status_code = 403
+        status_code = extract_http_status_code_tag(trigger_tags, response)
+        self.assertEqual(status_code, 403)