From 9bdaf018fd9e1c57f0c57146251b4f539a07ae4f Mon Sep 17 00:00:00 2001 From: Rey Abolofia Date: Fri, 5 Apr 2024 11:11:02 -0700 Subject: [PATCH] Move get_mock_context method to central utils file. --- tests/test_module_name.py | 2 -- tests/test_tags.py | 19 +++++++------------ tests/test_tracing.py | 28 ++++------------------------ tests/test_trigger.py | 8 ++------ tests/test_wrapper.py | 32 +++++++++----------------------- tests/utils.py | 26 ++++++++++++++++++++++++++ 6 files changed, 48 insertions(+), 67 deletions(-) create mode 100644 tests/utils.py diff --git a/tests/test_module_name.py b/tests/test_module_name.py index 21396465..a6faf829 100644 --- a/tests/test_module_name.py +++ b/tests/test_module_name.py @@ -1,7 +1,5 @@ import unittest -from unittest.mock import patch, call, MagicMock - from datadog_lambda.module_name import modify_module_name diff --git a/tests/test_tags.py b/tests/test_tags.py index 7a1ad86d..07daa8e0 100644 --- a/tests/test_tags.py +++ b/tests/test_tags.py @@ -1,19 +1,10 @@ import unittest -from unittest.mock import patch, MagicMock - +from unittest.mock import patch from datadog_lambda.tags import parse_lambda_tags_from_arn - -def get_mock_context( - invoked_function_arn="arn:aws:lambda:us-east-1:1234597598159:function:swf-hello-test:$Latest", - function_version="1", -): - lambda_context = MagicMock() - lambda_context.invoked_function_arn = invoked_function_arn - lambda_context.function_version = function_version - return lambda_context +from tests.utils import get_mock_context class TestMetricTags(unittest.TestCase): @@ -23,8 +14,12 @@ def setUp(self): self.addCleanup(patcher.stop) def test_parse_lambda_tags_from_arn_latest(self): + lambda_context = get_mock_context() + lambda_context.invoked_function_arn = ( + "arn:aws:lambda:us-east-1:1234597598159:function:swf-hello-test:$Latest" + ) self.assertListEqual( - parse_lambda_tags_from_arn(get_mock_context()), + parse_lambda_tags_from_arn(lambda_context), [ "region:us-east-1", "account_id:1234597598159", diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 0f7b84b4..296bd0dc 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -5,7 +5,7 @@ import os import unittest -from unittest.mock import MagicMock, Mock, patch, call +from unittest.mock import Mock, patch, call import ddtrace @@ -39,6 +39,9 @@ ) from datadog_lambda.trigger import EventTypes +from tests.utils import get_mock_context + + function_arn = "arn:aws:lambda:us-west-1:123457598159:function:python-layer-test" fake_xray_header_value = ( @@ -50,29 +53,6 @@ event_samples = "tests/event_samples/" -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", - function_name="Function", - 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.function_name = function_name - lambda_context.client_context = ClientContext(custom) - return lambda_context - - def with_trace_propagation_style(style): style_list = list(style.split(",")) diff --git a/tests/test_trigger.py b/tests/test_trigger.py index 2f514811..59178b1e 100644 --- a/tests/test_trigger.py +++ b/tests/test_trigger.py @@ -11,16 +11,12 @@ extract_http_status_code_tag, ) +from tests.utils import get_mock_context + event_samples = "tests/event_samples/" function_arn = "arn:aws:lambda:us-west-1:123457598159:function:python-layer-test" -def get_mock_context(invoked_function_arn=function_arn): - lambda_context = MagicMock() - lambda_context.invoked_function_arn = invoked_function_arn - return lambda_context - - class TestGetEventSourceAndARN(unittest.TestCase): def test_event_source_api_gateway(self): event_sample_source = "api-gateway" diff --git a/tests/test_wrapper.py b/tests/test_wrapper.py index ff038c21..b7058e11 100644 --- a/tests/test_wrapper.py +++ b/tests/test_wrapper.py @@ -3,7 +3,7 @@ import os import unittest -from unittest.mock import patch, call, ANY, MagicMock +from unittest.mock import patch, call, ANY from datadog_lambda.constants import TraceHeader import datadog_lambda.wrapper as wrapper @@ -12,21 +12,7 @@ from ddtrace import Span, tracer from ddtrace.internal.constants import MAX_UINT_64BITS - -def get_mock_context( - aws_request_id="request-id-1", - memory_limit_in_mb="256", - invoked_function_arn="arn:aws:lambda:us-west-1:123457598159:function:python-layer-test:1", - function_version="1", - client_context={}, -): - 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 - return lambda_context +from tests.utils import get_mock_context class TestDatadogLambdaWrapper(unittest.TestCase): @@ -232,7 +218,7 @@ def lambda_handler(event, context): "region:us-west-1", "account_id:123457598159", "functionname:python-layer-test", - "resource:python-layer-test:1", + "resource:python-layer-test", "cold_start:true", "memorysize:256", "runtime:python3.9", @@ -263,7 +249,7 @@ def lambda_handler(event, context): "region:us-west-1", "account_id:123457598159", "functionname:python-layer-test", - "resource:python-layer-test:1", + "resource:python-layer-test", "cold_start:true", "memorysize:256", "runtime:python3.9", @@ -279,7 +265,7 @@ def lambda_handler(event, context): "region:us-west-1", "account_id:123457598159", "functionname:python-layer-test", - "resource:python-layer-test:1", + "resource:python-layer-test", "cold_start:true", "memorysize:256", "runtime:python3.9", @@ -318,7 +304,7 @@ def lambda_handler(event, context): "region:us-west-1", "account_id:123457598159", "functionname:python-layer-test", - "resource:python-layer-test:1", + "resource:python-layer-test", "cold_start:true", "memorysize:256", "runtime:python3.9", @@ -334,7 +320,7 @@ def lambda_handler(event, context): "region:us-west-1", "account_id:123457598159", "functionname:python-layer-test", - "resource:python-layer-test:1", + "resource:python-layer-test", "cold_start:true", "memorysize:256", "runtime:python3.9", @@ -370,7 +356,7 @@ def lambda_handler(event, context): "region:us-west-1", "account_id:123457598159", "functionname:python-layer-test", - "resource:python-layer-test:1", + "resource:python-layer-test", "cold_start:true", "memorysize:256", "runtime:python3.9", @@ -386,7 +372,7 @@ def lambda_handler(event, context): "region:us-west-1", "account_id:123457598159", "functionname:python-layer-test", - "resource:python-layer-test:1", + "resource:python-layer-test", "cold_start:false", "memorysize:256", "runtime:python3.9", diff --git a/tests/utils.py b/tests/utils.py new file mode 100644 index 00000000..71060d46 --- /dev/null +++ b/tests/utils.py @@ -0,0 +1,26 @@ +from unittest.mock import MagicMock + +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", + function_name="Function", + 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.function_name = function_name + lambda_context.client_context = ClientContext(custom) + return lambda_context