forked from aws-powertools/powertools-lambda-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_tracing.py
More file actions
155 lines (112 loc) · 5.05 KB
/
Copy pathtest_tracing.py
File metadata and controls
155 lines (112 loc) · 5.05 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
from unittest import mock
import pytest
from aws_lambda_powertools.tracing import Tracer
@pytest.fixture
def dummy_response():
return {"test": "succeeds"}
@pytest.fixture
def xray_stub(mocker):
class XRayStub:
def __init__(
self,
put_metadata_mock: mocker.MagicMock = None,
put_annotation_mock: mocker.MagicMock = None,
begin_subsegment_mock: mocker.MagicMock = None,
end_subsegment_mock: mocker.MagicMock = None,
):
self.put_metadata_mock = put_metadata_mock or mocker.MagicMock()
self.put_annotation_mock = put_annotation_mock or mocker.MagicMock()
self.begin_subsegment_mock = begin_subsegment_mock or mocker.MagicMock()
self.end_subsegment_mock = end_subsegment_mock or mocker.MagicMock()
def put_metadata(self, *args, **kwargs):
return self.put_metadata_mock(*args, **kwargs)
def put_annotation(self, *args, **kwargs):
return self.put_annotation_mock(*args, **kwargs)
def begin_subsegment(self, *args, **kwargs):
return self.begin_subsegment_mock(*args, **kwargs)
def end_subsegment(self, *args, **kwargs):
return self.end_subsegment_mock(*args, **kwargs)
return XRayStub
@pytest.fixture(scope="function", autouse=True)
def reset_tracing_config():
Tracer._reset_config()
yield
def test_tracer_lambda_handler(mocker, dummy_response, xray_stub):
put_metadata_mock = mocker.MagicMock()
begin_subsegment_mock = mocker.MagicMock()
end_subsegment_mock = mocker.MagicMock()
xray_provider = xray_stub(
put_metadata_mock=put_metadata_mock,
begin_subsegment_mock=begin_subsegment_mock,
end_subsegment_mock=end_subsegment_mock,
)
tracer = Tracer(provider=xray_provider, service="booking")
@tracer.capture_lambda_handler
def handler(event, context):
return dummy_response
handler({}, mocker.MagicMock())
assert begin_subsegment_mock.call_count == 1
assert begin_subsegment_mock.call_args == mocker.call(name="## handler")
assert end_subsegment_mock.call_count == 1
assert put_metadata_mock.call_args == mocker.call(
key="lambda handler response", value=dummy_response, namespace="booking"
)
def test_tracer_method(mocker, dummy_response, xray_stub):
put_metadata_mock = mocker.MagicMock()
put_annotation_mock = mocker.MagicMock()
begin_subsegment_mock = mocker.MagicMock()
end_subsegment_mock = mocker.MagicMock()
xray_provider = xray_stub(put_metadata_mock, put_annotation_mock, begin_subsegment_mock, end_subsegment_mock)
tracer = Tracer(provider=xray_provider, service="booking")
@tracer.capture_method
def greeting(name, message):
return dummy_response
greeting(name="Foo", message="Bar")
assert begin_subsegment_mock.call_count == 1
assert begin_subsegment_mock.call_args == mocker.call(name="## greeting")
assert end_subsegment_mock.call_count == 1
assert put_metadata_mock.call_args == mocker.call(
key="greeting response", value=dummy_response, namespace="booking"
)
def test_tracer_custom_metadata(mocker, dummy_response, xray_stub):
put_metadata_mock = mocker.MagicMock()
xray_provider = xray_stub(put_metadata_mock=put_metadata_mock)
tracer = Tracer(provider=xray_provider, service="booking")
annotation_key = "Booking response"
annotation_value = {"bookingStatus": "CONFIRMED"}
@tracer.capture_lambda_handler
def handler(event, context):
tracer.put_metadata(annotation_key, annotation_value)
return dummy_response
handler({}, mocker.MagicMock())
assert put_metadata_mock.call_count == 2
assert put_metadata_mock.call_args_list[0] == mocker.call(
key=annotation_key, value=annotation_value, namespace="booking"
)
def test_tracer_custom_annotation(mocker, dummy_response, xray_stub):
put_annotation_mock = mocker.MagicMock()
xray_provider = xray_stub(put_annotation_mock=put_annotation_mock)
tracer = Tracer(provider=xray_provider, service="booking")
annotation_key = "BookingId"
annotation_value = "123456"
@tracer.capture_lambda_handler
def handler(event, context):
tracer.put_annotation(annotation_key, annotation_value)
return dummy_response
handler({}, mocker.MagicMock())
assert put_annotation_mock.call_count == 1
assert put_annotation_mock.call_args == mocker.call(key=annotation_key, value=annotation_value)
@mock.patch("aws_lambda_powertools.tracing.Tracer.patch")
def test_tracer_autopatch(patch_mock):
# GIVEN tracer is instantiated
# WHEN default options were used, or patch() was called
# THEN tracer should patch all modules
Tracer(disabled=True)
assert patch_mock.call_count == 1
@mock.patch("aws_lambda_powertools.tracing.Tracer.patch")
def test_tracer_no_autopatch(patch_mock):
# GIVEN tracer is instantiated
# WHEN auto_patch is disabled
# THEN tracer should not patch any module
Tracer(disabled=True, auto_patch=False)
assert patch_mock.call_count == 0