Skip to content

Commit cfeed17

Browse files
purple4reinapurple4reina
authored
Move get_mock_context method to central utils file. (#465)
1 parent 36610c1 commit cfeed17

File tree

6 files changed

+48
-67
lines changed

6 files changed

+48
-67
lines changed

‎tests/test_module_name.py

-2
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,5 @@
11
import unittest
22

3-
from unittest.mock import patch, call, MagicMock
4-
53
from datadog_lambda.module_name import modify_module_name
64

75

‎tests/test_tags.py

+7-12
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,10 @@
11
import unittest
22

3-
from unittest.mock import patch, MagicMock
4-
3+
from unittest.mock import patch
54

65
from datadog_lambda.tags import parse_lambda_tags_from_arn
76

8-
9-
def get_mock_context(
10-
invoked_function_arn="arn:aws:lambda:us-east-1:1234597598159:function:swf-hello-test:$Latest",
11-
function_version="1",
12-
):
13-
lambda_context = MagicMock()
14-
lambda_context.invoked_function_arn = invoked_function_arn
15-
lambda_context.function_version = function_version
16-
return lambda_context
7+
from tests.utils import get_mock_context
178

189

1910
class TestMetricTags(unittest.TestCase):
@@ -23,8 +14,12 @@ def setUp(self):
2314
self.addCleanup(patcher.stop)
2415

2516
def test_parse_lambda_tags_from_arn_latest(self):
17+
lambda_context = get_mock_context()
18+
lambda_context.invoked_function_arn = (
19+
"arn:aws:lambda:us-east-1:1234597598159:function:swf-hello-test:$Latest"
20+
)
2621
self.assertListEqual(
27-
parse_lambda_tags_from_arn(get_mock_context()),
22+
parse_lambda_tags_from_arn(lambda_context),
2823
[
2924
"region:us-east-1",
3025
"account_id:1234597598159",

‎tests/test_tracing.py

+4-24
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
import os
66
import unittest
77

8-
from unittest.mock import MagicMock, Mock, patch, call
8+
from unittest.mock import Mock, patch, call
99

1010
import ddtrace
1111

@@ -39,6 +39,9 @@
3939
)
4040
from datadog_lambda.trigger import EventTypes
4141

42+
from tests.utils import get_mock_context
43+
44+
4245
function_arn = "arn:aws:lambda:us-west-1:123457598159:function:python-layer-test"
4346

4447
fake_xray_header_value = (
@@ -50,29 +53,6 @@
5053
event_samples = "tests/event_samples/"
5154

5255

53-
class ClientContext(object):
54-
def __init__(self, custom=None):
55-
self.custom = custom
56-
57-
58-
def get_mock_context(
59-
aws_request_id="request-id-1",
60-
memory_limit_in_mb="256",
61-
invoked_function_arn=function_arn,
62-
function_version="1",
63-
function_name="Function",
64-
custom=None,
65-
):
66-
lambda_context = MagicMock()
67-
lambda_context.aws_request_id = aws_request_id
68-
lambda_context.memory_limit_in_mb = memory_limit_in_mb
69-
lambda_context.invoked_function_arn = invoked_function_arn
70-
lambda_context.function_version = function_version
71-
lambda_context.function_name = function_name
72-
lambda_context.client_context = ClientContext(custom)
73-
return lambda_context
74-
75-
7656
def with_trace_propagation_style(style):
7757
style_list = list(style.split(","))
7858

‎tests/test_trigger.py

+2-6
Original file line numberDiff line numberDiff line change
@@ -11,16 +11,12 @@
1111
extract_http_status_code_tag,
1212
)
1313

14+
from tests.utils import get_mock_context
15+
1416
event_samples = "tests/event_samples/"
1517
function_arn = "arn:aws:lambda:us-west-1:123457598159:function:python-layer-test"
1618

1719

18-
def get_mock_context(invoked_function_arn=function_arn):
19-
lambda_context = MagicMock()
20-
lambda_context.invoked_function_arn = invoked_function_arn
21-
return lambda_context
22-
23-
2420
class TestGetEventSourceAndARN(unittest.TestCase):
2521
def test_event_source_api_gateway(self):
2622
event_sample_source = "api-gateway"

‎tests/test_wrapper.py

+9-23
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
import os
44
import unittest
55

6-
from unittest.mock import patch, call, ANY, MagicMock
6+
from unittest.mock import patch, call, ANY
77
from datadog_lambda.constants import TraceHeader
88

99
import datadog_lambda.wrapper as wrapper
@@ -12,21 +12,7 @@
1212
from ddtrace import Span, tracer
1313
from ddtrace.internal.constants import MAX_UINT_64BITS
1414

15-
16-
def get_mock_context(
17-
aws_request_id="request-id-1",
18-
memory_limit_in_mb="256",
19-
invoked_function_arn="arn:aws:lambda:us-west-1:123457598159:function:python-layer-test:1",
20-
function_version="1",
21-
client_context={},
22-
):
23-
lambda_context = MagicMock()
24-
lambda_context.aws_request_id = aws_request_id
25-
lambda_context.memory_limit_in_mb = memory_limit_in_mb
26-
lambda_context.invoked_function_arn = invoked_function_arn
27-
lambda_context.function_version = function_version
28-
lambda_context.client_context = client_context
29-
return lambda_context
15+
from tests.utils import get_mock_context
3016

3117

3218
class TestDatadogLambdaWrapper(unittest.TestCase):
@@ -232,7 +218,7 @@ def lambda_handler(event, context):
232218
"region:us-west-1",
233219
"account_id:123457598159",
234220
"functionname:python-layer-test",
235-
"resource:python-layer-test:1",
221+
"resource:python-layer-test",
236222
"cold_start:true",
237223
"memorysize:256",
238224
"runtime:python3.9",
@@ -263,7 +249,7 @@ def lambda_handler(event, context):
263249
"region:us-west-1",
264250
"account_id:123457598159",
265251
"functionname:python-layer-test",
266-
"resource:python-layer-test:1",
252+
"resource:python-layer-test",
267253
"cold_start:true",
268254
"memorysize:256",
269255
"runtime:python3.9",
@@ -279,7 +265,7 @@ def lambda_handler(event, context):
279265
"region:us-west-1",
280266
"account_id:123457598159",
281267
"functionname:python-layer-test",
282-
"resource:python-layer-test:1",
268+
"resource:python-layer-test",
283269
"cold_start:true",
284270
"memorysize:256",
285271
"runtime:python3.9",
@@ -318,7 +304,7 @@ def lambda_handler(event, context):
318304
"region:us-west-1",
319305
"account_id:123457598159",
320306
"functionname:python-layer-test",
321-
"resource:python-layer-test:1",
307+
"resource:python-layer-test",
322308
"cold_start:true",
323309
"memorysize:256",
324310
"runtime:python3.9",
@@ -334,7 +320,7 @@ def lambda_handler(event, context):
334320
"region:us-west-1",
335321
"account_id:123457598159",
336322
"functionname:python-layer-test",
337-
"resource:python-layer-test:1",
323+
"resource:python-layer-test",
338324
"cold_start:true",
339325
"memorysize:256",
340326
"runtime:python3.9",
@@ -370,7 +356,7 @@ def lambda_handler(event, context):
370356
"region:us-west-1",
371357
"account_id:123457598159",
372358
"functionname:python-layer-test",
373-
"resource:python-layer-test:1",
359+
"resource:python-layer-test",
374360
"cold_start:true",
375361
"memorysize:256",
376362
"runtime:python3.9",
@@ -386,7 +372,7 @@ def lambda_handler(event, context):
386372
"region:us-west-1",
387373
"account_id:123457598159",
388374
"functionname:python-layer-test",
389-
"resource:python-layer-test:1",
375+
"resource:python-layer-test",
390376
"cold_start:false",
391377
"memorysize:256",
392378
"runtime:python3.9",

‎tests/utils.py

+26
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
from unittest.mock import MagicMock
2+
3+
function_arn = "arn:aws:lambda:us-west-1:123457598159:function:python-layer-test"
4+
5+
6+
class ClientContext(object):
7+
def __init__(self, custom=None):
8+
self.custom = custom
9+
10+
11+
def get_mock_context(
12+
aws_request_id="request-id-1",
13+
memory_limit_in_mb="256",
14+
invoked_function_arn=function_arn,
15+
function_version="1",
16+
function_name="Function",
17+
custom=None,
18+
):
19+
lambda_context = MagicMock()
20+
lambda_context.aws_request_id = aws_request_id
21+
lambda_context.memory_limit_in_mb = memory_limit_in_mb
22+
lambda_context.invoked_function_arn = invoked_function_arn
23+
lambda_context.function_version = function_version
24+
lambda_context.function_name = function_name
25+
lambda_context.client_context = ClientContext(custom)
26+
return lambda_context

0 commit comments

Comments
 (0)