Move get_mock_context method to central utils file.

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
 
 

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",

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(","))
 

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"

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",

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