From 70148fda020f3e337aa9caa43bc1d425e53c6bc3 Mon Sep 17 00:00:00 2001 From: Tian Chu Date: Tue, 9 Mar 2021 22:31:56 -0500 Subject: [PATCH] Fix get status_code --- datadog_lambda/trigger.py | 4 +++- tests/test_trigger.py | 16 ++++++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/datadog_lambda/trigger.py b/datadog_lambda/trigger.py index b4427eed..b574b8f4 100644 --- a/datadog_lambda/trigger.py +++ b/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 diff --git a/tests/test_trigger.py b/tests/test_trigger.py index 34b7532c..bc60e85a 100644 --- a/tests/test_trigger.py +++ b/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)