Skip to content

Commit fc8b1fd

Browse files
DarcyRaynerDDDarcyRaynerDD
committed
Extend tracing tests
1 parent 082683c commit fc8b1fd

File tree

1 file changed

+98
-79
lines changed

1 file changed

+98
-79
lines changed

‎tests/test_tracing.py‎

Lines changed: 98 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -1,182 +1,201 @@
11
import unittest
2+
23
try:
34
from unittest.mock import MagicMock, patch
45
except ImportError:
56
from mock import MagicMock, patch
67

78
from ddtrace.helpers import get_correlation_ids
8-
from datadog_lambda.constants import (
9-
SamplingPriority,
10-
TraceHeader,
11-
XraySubsegment,
12-
)
9+
from ddtrace import tracer
10+
11+
from datadog_lambda.constants import SamplingPriority, TraceHeader, XraySubsegment
1312
from datadog_lambda.tracing import (
1413
extract_dd_trace_context,
1514
get_dd_trace_context,
1615
set_correlation_ids,
1716
_convert_xray_trace_id,
1817
_convert_xray_entity_id,
1918
_convert_xray_sampling,
19+
dd_native_tracing_enabled,
2020
)
2121

2222

2323
class TestExtractAndGetDDTraceContext(unittest.TestCase):
24-
2524
def setUp(self):
26-
patcher = patch('datadog_lambda.tracing.xray_recorder')
25+
global dd_native_tracing_enabled
26+
dd_native_tracing_enabled = False
27+
patcher = patch("datadog_lambda.tracing.xray_recorder")
2728
self.mock_xray_recorder = patcher.start()
2829
self.mock_xray_recorder.get_trace_entity.return_value = MagicMock(
29-
id='ffff',
30-
trace_id='1111',
31-
sampled=True,
30+
id="ffff", trace_id="1111", sampled=True
3231
)
3332
self.mock_current_subsegment = MagicMock()
34-
self.mock_xray_recorder.current_subsegment.return_value = \
33+
self.mock_xray_recorder.current_subsegment.return_value = (
3534
self.mock_current_subsegment
35+
)
3636
self.addCleanup(patcher.stop)
3737

38-
patcher = patch('datadog_lambda.tracing.is_lambda_context')
38+
patcher = patch("datadog_lambda.tracing.is_lambda_context")
3939
self.mock_is_lambda_context = patcher.start()
4040
self.mock_is_lambda_context.return_value = True
4141
self.addCleanup(patcher.stop)
4242

43+
def tearDown(self):
44+
global dd_native_tracing_enabled
45+
dd_native_tracing_enabled = False
46+
4347
def test_without_datadog_trace_headers(self):
44-
extract_dd_trace_context({})
48+
ctx = extract_dd_trace_context({})
49+
self.assertDictEqual(
50+
ctx,
51+
{
52+
"trace-id": "4369",
53+
"parent-id": "65535",
54+
"sampling-priority": "2",
55+
"source": "xray",
56+
},
57+
)
4558
self.assertDictEqual(
4659
get_dd_trace_context(),
4760
{
48-
TraceHeader.TRACE_ID: '4369',
49-
TraceHeader.PARENT_ID: '65535',
50-
TraceHeader.SAMPLING_PRIORITY: '2',
51-
}
61+
TraceHeader.TRACE_ID: "4369",
62+
TraceHeader.PARENT_ID: "65535",
63+
TraceHeader.SAMPLING_PRIORITY: "2",
64+
},
5265
)
5366

5467
def test_with_incomplete_datadog_trace_headers(self):
55-
extract_dd_trace_context({
56-
'headers': {
57-
TraceHeader.TRACE_ID: '123',
58-
TraceHeader.PARENT_ID: '321',
59-
}
60-
})
68+
ctx = extract_dd_trace_context(
69+
{"headers": {TraceHeader.TRACE_ID: "123", TraceHeader.PARENT_ID: "321"}}
70+
)
71+
self.assertDictEqual(
72+
ctx,
73+
{
74+
"trace-id": "4369",
75+
"parent-id": "65535",
76+
"sampling-priority": "2",
77+
"source": "xray",
78+
},
79+
)
6180
self.assertDictEqual(
6281
get_dd_trace_context(),
6382
{
64-
TraceHeader.TRACE_ID: '4369',
65-
TraceHeader.PARENT_ID: '65535',
66-
TraceHeader.SAMPLING_PRIORITY: '2',
67-
}
83+
TraceHeader.TRACE_ID: "4369",
84+
TraceHeader.PARENT_ID: "65535",
85+
TraceHeader.SAMPLING_PRIORITY: "2",
86+
},
6887
)
6988

7089
def test_with_complete_datadog_trace_headers(self):
71-
extract_dd_trace_context({
72-
'headers': {
73-
TraceHeader.TRACE_ID: '123',
74-
TraceHeader.PARENT_ID: '321',
75-
TraceHeader.SAMPLING_PRIORITY: '1',
90+
ctx = extract_dd_trace_context(
91+
{
92+
"headers": {
93+
TraceHeader.TRACE_ID: "123",
94+
TraceHeader.PARENT_ID: "321",
95+
TraceHeader.SAMPLING_PRIORITY: "1",
96+
}
7697
}
77-
})
98+
)
99+
self.assertDictEqual(
100+
ctx,
101+
{
102+
"trace-id": "123",
103+
"parent-id": "321",
104+
"sampling-priority": "1",
105+
"source": "event",
106+
},
107+
)
78108
self.assertDictEqual(
79109
get_dd_trace_context(),
80110
{
81-
TraceHeader.TRACE_ID: '123',
82-
TraceHeader.PARENT_ID: '65535',
83-
TraceHeader.SAMPLING_PRIORITY: '1',
84-
}
111+
TraceHeader.TRACE_ID: "123",
112+
TraceHeader.PARENT_ID: "65535",
113+
TraceHeader.SAMPLING_PRIORITY: "1",
114+
},
85115
)
86116
self.mock_xray_recorder.begin_subsegment.assert_called()
87117
self.mock_xray_recorder.end_subsegment.assert_called()
88118
self.mock_current_subsegment.put_metadata.assert_called_with(
89119
XraySubsegment.KEY,
90-
{
91-
'trace-id': '123',
92-
'parent-id': '321',
93-
'sampling-priority': '1',
94-
},
95-
XraySubsegment.NAMESPACE
120+
{"trace-id": "123", "parent-id": "321", "sampling-priority": "1"},
121+
XraySubsegment.NAMESPACE,
96122
)
97123

98124
def test_with_complete_datadog_trace_headers_with_mixed_casing(self):
99-
extract_dd_trace_context({
100-
'headers': {
101-
'X-Datadog-Trace-Id': '123',
102-
'X-Datadog-Parent-Id': '321',
103-
'X-Datadog-Sampling-Priority': '1',
125+
extract_dd_trace_context(
126+
{
127+
"headers": {
128+
"X-Datadog-Trace-Id": "123",
129+
"X-Datadog-Parent-Id": "321",
130+
"X-Datadog-Sampling-Priority": "1",
131+
}
104132
}
105-
})
133+
)
106134
self.assertDictEqual(
107135
get_dd_trace_context(),
108136
{
109-
TraceHeader.TRACE_ID: '123',
110-
TraceHeader.PARENT_ID: '65535',
111-
TraceHeader.SAMPLING_PRIORITY: '1',
112-
}
137+
TraceHeader.TRACE_ID: "123",
138+
TraceHeader.PARENT_ID: "65535",
139+
TraceHeader.SAMPLING_PRIORITY: "1",
140+
},
113141
)
114142

115143

116144
class TestXRayContextConversion(unittest.TestCase):
117-
118145
def test_convert_xray_trace_id(self):
119146
self.assertEqual(
120-
_convert_xray_trace_id('00000000e1be46a994272793'),
121-
'7043144561403045779'
147+
_convert_xray_trace_id("00000000e1be46a994272793"), "7043144561403045779"
122148
)
123149

124150
self.assertEqual(
125-
_convert_xray_trace_id('bd862e3fe1be46a994272793'),
126-
'7043144561403045779'
151+
_convert_xray_trace_id("bd862e3fe1be46a994272793"), "7043144561403045779"
127152
)
128153

129154
self.assertEqual(
130-
_convert_xray_trace_id('ffffffffffffffffffffffff'),
131-
'9223372036854775807' # 0x7FFFFFFFFFFFFFFF
155+
_convert_xray_trace_id("ffffffffffffffffffffffff"),
156+
"9223372036854775807", # 0x7FFFFFFFFFFFFFFF
132157
)
133158

134159
def test_convert_xray_entity_id(self):
135160
self.assertEqual(
136-
_convert_xray_entity_id('53995c3f42cd8ad8'),
137-
'6023947403358210776'
161+
_convert_xray_entity_id("53995c3f42cd8ad8"), "6023947403358210776"
138162
)
139163

140164
self.assertEqual(
141-
_convert_xray_entity_id('1000000000000000'),
142-
'1152921504606846976'
165+
_convert_xray_entity_id("1000000000000000"), "1152921504606846976"
143166
)
144167

145168
self.assertEqual(
146-
_convert_xray_entity_id('ffffffffffffffff'),
147-
'18446744073709551615'
169+
_convert_xray_entity_id("ffffffffffffffff"), "18446744073709551615"
148170
)
149171

150172
def test_convert_xray_sampling(self):
151-
self.assertEqual(
152-
_convert_xray_sampling(True),
153-
str(SamplingPriority.USER_KEEP)
154-
)
173+
self.assertEqual(_convert_xray_sampling(True), str(SamplingPriority.USER_KEEP))
155174

156175
self.assertEqual(
157-
_convert_xray_sampling(False),
158-
str(SamplingPriority.USER_REJECT)
176+
_convert_xray_sampling(False), str(SamplingPriority.USER_REJECT)
159177
)
160178

161179

162180
class TestLogsInjection(unittest.TestCase):
163-
164181
def setUp(self):
165-
patcher = patch('datadog_lambda.tracing.get_dd_trace_context')
182+
183+
patcher = patch("datadog_lambda.tracing.get_dd_trace_context")
166184
self.mock_get_dd_trace_context = patcher.start()
167185
self.mock_get_dd_trace_context.return_value = {
168-
TraceHeader.TRACE_ID: '123',
169-
TraceHeader.PARENT_ID: '456',
186+
TraceHeader.TRACE_ID: "123",
187+
TraceHeader.PARENT_ID: "456",
170188
}
171189
self.addCleanup(patcher.stop)
172190

173-
patcher = patch('datadog_lambda.tracing.is_lambda_context')
191+
patcher = patch("datadog_lambda.tracing.is_lambda_context")
174192
self.mock_is_lambda_context = patcher.start()
175193
self.mock_is_lambda_context.return_value = True
176194
self.addCleanup(patcher.stop)
177195

178196
def test_set_correlation_ids(self):
179197
set_correlation_ids()
180198
trace_id, span_id = get_correlation_ids()
181-
self.assertEqual(trace_id, '123')
182-
self.assertEqual(span_id, '456')
199+
self.assertEqual(trace_id, "123")
200+
self.assertEqual(span_id, "456")
201+

0 commit comments

Comments
 (0)