Skip to content

Commit 77af3c7

Browse files
michael-zhao459michael-zhao459
committed
fixed lint for tests
1 parent a9c54be commit 77af3c7

File tree

1 file changed

+49
-42
lines changed

1 file changed

+49
-42
lines changed

‎tests/test_dsm.py‎

Lines changed: 49 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ def setUp(self):
3030
self.addCleanup(logger_patcher.stop)
3131

3232
def test_extract_context_data_streams_disabled(self):
33-
with patch.dict(os.environ, {'DD_DATA_STREAMS_ENABLED': 'false'}):
33+
with patch.dict(os.environ, {"DD_DATA_STREAMS_ENABLED": "false"}):
3434
context_json = {"dd-pathway-ctx-base64": "12345"}
3535
event_type = "sqs"
3636
arn = "arn:aws:sqs:us-east-1:123456789012:test-queue"
@@ -45,7 +45,7 @@ def test_extract_context_data_streams_disabled(self):
4545
self.assertEqual(result, mock_context)
4646

4747
def test_extract_context_data_streams_enabled_complete_context(self):
48-
with patch.dict(os.environ, {'DD_DATA_STREAMS_ENABLED': 'true'}):
48+
with patch.dict(os.environ, {"DD_DATA_STREAMS_ENABLED": "true"}):
4949
context_json = {"dd-pathway-ctx-base64": "12345"}
5050
event_type = "sqs"
5151
arn = "arn:aws:sqs:us-east-1:123456789012:test-queue"
@@ -65,7 +65,7 @@ def test_extract_context_data_streams_enabled_complete_context(self):
6565
self.assertEqual(result, mock_context)
6666

6767
def test_extract_context_data_streams_enabled_incomplete_context(self):
68-
with patch.dict(os.environ, {'DD_DATA_STREAMS_ENABLED': 'true'}):
68+
with patch.dict(os.environ, {"DD_DATA_STREAMS_ENABLED": "true"}):
6969
context_json = {"dd-pathway-ctx-base64": "12345"}
7070
event_type = "sqs"
7171
arn = "arn:aws:sqs:us-east-1:123456789012:test-queue"
@@ -80,7 +80,7 @@ def test_extract_context_data_streams_enabled_incomplete_context(self):
8080
self.assertEqual(result, mock_context)
8181

8282
def test_extract_context_exception_path(self):
83-
with patch.dict(os.environ, {'DD_DATA_STREAMS_ENABLED': 'true'}):
83+
with patch.dict(os.environ, {"DD_DATA_STREAMS_ENABLED": "true"}):
8484
context_json = {"dd-pathway-ctx-base64": "12345"}
8585
event_type = "sqs"
8686
arn = "arn:aws:sqs:us-east-1:123456789012:test-queue"
@@ -104,7 +104,7 @@ def test_create_carrier_get_with_valid_data(self):
104104
context_json = {
105105
"x-datadog-trace-id": "12345",
106106
"x-datadog-parent-id": "67890",
107-
"x-datadog-sampling-priority": "1"
107+
"x-datadog-sampling-priority": "1",
108108
}
109109

110110
carrier_get = _create_carrier_get(context_json)
@@ -142,49 +142,53 @@ def test_sqs_event_with_datadog_message_attributes(self, mock_extract_context):
142142
dd_json_data = json.dumps(dd_data)
143143

144144
event = {
145-
"Records": [{
146-
"eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue",
147-
"messageAttributes": {
148-
"_datadog": {
149-
"dataType": "String",
150-
"stringValue": dd_json_data
151-
}
145+
"Records": [
146+
{
147+
"eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue",
148+
"messageAttributes": {
149+
"_datadog": {"dataType": "String", "stringValue": dd_json_data}
150+
},
152151
}
153-
}]
152+
]
154153
}
155154

156155
mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1)
157156
mock_extract_context.return_value = mock_context
158157

159-
result = extract_context_from_sqs_or_sns_event_or_context(event, self.lambda_context)
158+
result = extract_context_from_sqs_or_sns_event_or_context(
159+
event, self.lambda_context
160+
)
160161

161162
mock_extract_context.assert_called_once_with(
162163
dd_data, "sqs", "arn:aws:sqs:us-east-1:123456789012:test-queue"
163164
)
164165
self.assertEqual(result, mock_context)
165166

166167
@patch("datadog_lambda.tracing._extract_context")
167-
def test_sqs_event_with_binary_datadog_message_attributes(self, mock_extract_context):
168+
def test_sqs_event_with_binary_datadog_message_attributes(
169+
self, mock_extract_context
170+
):
168171
dd_data = {"dd-pathway-ctx-base64": "12345"}
169172
dd_json_data = json.dumps(dd_data)
170173
encoded_data = base64.b64encode(dd_json_data.encode()).decode()
171174

172175
event = {
173-
"Records": [{
174-
"eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue",
175-
"messageAttributes": {
176-
"_datadog": {
177-
"dataType": "Binary",
178-
"binaryValue": encoded_data
179-
}
176+
"Records": [
177+
{
178+
"eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue",
179+
"messageAttributes": {
180+
"_datadog": {"dataType": "Binary", "binaryValue": encoded_data}
181+
},
180182
}
181-
}]
183+
]
182184
}
183185

184186
mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1)
185187
mock_extract_context.return_value = mock_context
186188

187-
result = extract_context_from_sqs_or_sns_event_or_context(event, self.lambda_context)
189+
result = extract_context_from_sqs_or_sns_event_or_context(
190+
event, self.lambda_context
191+
)
188192

189193
mock_extract_context.assert_called_once_with(
190194
dd_data, "sqs", "arn:aws:sqs:us-east-1:123456789012:test-queue"
@@ -197,24 +201,25 @@ def test_sns_event_with_datadog_message_attributes(self, mock_extract_context):
197201
dd_json_data = json.dumps(dd_data)
198202

199203
event = {
200-
"Records": [{
201-
"eventSourceARN": "",
202-
"Sns": {
203-
"TopicArn": "arn:aws:sns:us-east-1:123456789012:test-topic",
204-
"MessageAttributes": {
205-
"_datadog": {
206-
"Type": "String",
207-
"Value": dd_json_data
208-
}
209-
}
204+
"Records": [
205+
{
206+
"eventSourceARN": "",
207+
"Sns": {
208+
"TopicArn": "arn:aws:sns:us-east-1:123456789012:test-topic",
209+
"MessageAttributes": {
210+
"_datadog": {"Type": "String", "Value": dd_json_data}
211+
},
212+
},
210213
}
211-
}]
214+
]
212215
}
213216

214217
mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1)
215218
mock_extract_context.return_value = mock_context
216219

217-
result = extract_context_from_sqs_or_sns_event_or_context(event, self.lambda_context)
220+
result = extract_context_from_sqs_or_sns_event_or_context(
221+
event, self.lambda_context
222+
)
218223

219224
mock_extract_context.assert_called_once_with(
220225
dd_data, "sns", "arn:aws:sns:us-east-1:123456789012:test-topic"
@@ -234,12 +239,12 @@ def test_kinesis_event_with_datadog_data(self, mock_extract_context):
234239
encoded_data = base64.b64encode(kinesis_data_str.encode()).decode()
235240

236241
event = {
237-
"Records": [{
238-
"eventSourceARN": "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream",
239-
"kinesis": {
240-
"data": encoded_data
242+
"Records": [
243+
{
244+
"eventSourceARN": "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream",
245+
"kinesis": {"data": encoded_data},
241246
}
242-
}]
247+
]
243248
}
244249

245250
mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1)
@@ -248,6 +253,8 @@ def test_kinesis_event_with_datadog_data(self, mock_extract_context):
248253
result = extract_context_from_kinesis_event(event, self.lambda_context)
249254

250255
mock_extract_context.assert_called_once_with(
251-
dd_data, "kinesis", "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream"
256+
dd_data,
257+
"kinesis",
258+
"arn:aws:kinesis:us-east-1:123456789012:stream/test-stream",
252259
)
253260
self.assertEqual(result, mock_context)

0 commit comments

Comments
 (0)