Skip to content

Commit dda744e

Browse files
michael-zhao459michael-zhao459
committed
unit test fixes
1 parent 254b864 commit dda744e

File tree

2 files changed

+97
-98
lines changed

2 files changed

+97
-98
lines changed

‎tests/test_dsm.py‎

Lines changed: 51 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -6,18 +6,42 @@
66

77

88
class TestDsmContext(unittest.TestCase):
9+
def setUp(self):
10+
patcher = patch("datadog_lambda.dsm._dsm_set_sqs_context")
11+
self.mock_dsm_set_sqs_context = patcher.start()
12+
self.addCleanup(patcher.stop)
13+
14+
patcher = patch("ddtrace.internal.datastreams.data_streams_processor")
15+
self.mock_data_streams_processor = patcher.start()
16+
self.addCleanup(patcher.stop)
17+
18+
patcher = patch("ddtrace.internal.datastreams.botocore.get_datastreams_context")
19+
self.mock_get_datastreams_context = patcher.start()
20+
self.mock_get_datastreams_context.return_value = {}
21+
self.addCleanup(patcher.stop)
22+
23+
patcher = patch(
24+
"ddtrace.internal.datastreams.botocore.calculate_sqs_payload_size"
25+
)
26+
self.mock_calculate_sqs_payload_size = patcher.start()
27+
self.mock_calculate_sqs_payload_size.return_value = 100
28+
self.addCleanup(patcher.stop)
29+
30+
patcher = patch("ddtrace.internal.datastreams.processor.DsmPathwayCodec.decode")
31+
self.mock_dsm_pathway_codec_decode = patcher.start()
32+
self.addCleanup(patcher.stop)
33+
934
def test_non_sqs_event_source_does_nothing(self):
1035
"""Test that non-SQS event sources don't trigger DSM context setting"""
1136
event = {"Records": [{"body": "test"}]}
1237

1338
mock_event_source = MagicMock()
1439
mock_event_source.equals.return_value = False # Not SQS
1540

16-
with patch("datadog_lambda.dsm._dsm_set_sqs_context") as mock_sqs_context:
17-
set_dsm_context(event, mock_event_source)
41+
set_dsm_context(event, mock_event_source)
1842

19-
mock_event_source.equals.assert_called_once_with(EventTypes.SQS)
20-
mock_sqs_context.assert_not_called()
43+
mock_event_source.equals.assert_called_once_with(EventTypes.SQS)
44+
self.mock_dsm_set_sqs_context.assert_not_called()
2145

2246
def test_event_with_no_records_does_nothing(self):
2347
"""Test that events where Records is None don't trigger DSM processing"""
@@ -28,12 +52,8 @@ def test_event_with_no_records_does_nothing(self):
2852
]
2953

3054
for event in events_with_no_records:
31-
with patch(
32-
"ddtrace.internal.datastreams.data_streams_processor"
33-
) as mock_processor:
34-
_dsm_set_sqs_context(event)
35-
36-
mock_processor.assert_not_called()
55+
_dsm_set_sqs_context(event)
56+
self.mock_data_streams_processor.assert_not_called()
3757

3858
def test_sqs_event_triggers_dsm_sqs_context(self):
3959
"""Test that SQS event sources trigger the SQS-specific DSM context function"""
@@ -50,10 +70,9 @@ def test_sqs_event_triggers_dsm_sqs_context(self):
5070
mock_event_source = MagicMock()
5171
mock_event_source.equals.return_value = True
5272

53-
with patch("datadog_lambda.dsm._dsm_set_sqs_context") as mock_sqs_context:
54-
set_dsm_context(sqs_event, mock_event_source)
73+
set_dsm_context(sqs_event, mock_event_source)
5574

56-
mock_sqs_context.assert_called_once_with(sqs_event)
75+
self.mock_dsm_set_sqs_context.assert_called_once_with(sqs_event)
5776

5877
def test_multiple_records_process_each_record(self):
5978
"""Test that each record in an SQS event gets processed individually"""
@@ -74,40 +93,24 @@ def test_multiple_records_process_each_record(self):
7493
]
7594
}
7695

77-
mock_processor = MagicMock()
7896
mock_context = MagicMock()
97+
self.mock_dsm_pathway_codec_decode.return_value = mock_context
98+
99+
_dsm_set_sqs_context(multi_record_event)
100+
101+
self.assertEqual(mock_context.set_checkpoint.call_count, 3)
102+
103+
calls = mock_context.set_checkpoint.call_args_list
104+
expected_arns = [
105+
"arn:aws:sqs:us-east-1:123456789012:queue1",
106+
"arn:aws:sqs:us-east-1:123456789012:queue2",
107+
"arn:aws:sqs:us-east-1:123456789012:queue3",
108+
]
79109

80-
with patch(
81-
"ddtrace.internal.datastreams.data_streams_processor",
82-
return_value=mock_processor,
83-
):
84-
with patch(
85-
"ddtrace.internal.datastreams.botocore.get_datastreams_context",
86-
return_value={},
87-
):
88-
with patch(
89-
"ddtrace.internal.datastreams.botocore.calculate_sqs_payload_size",
90-
return_value=100,
91-
):
92-
with patch(
93-
"ddtrace.internal.datastreams.processor.DsmPathwayCodec.decode",
94-
return_value=mock_context,
95-
):
96-
_dsm_set_sqs_context(multi_record_event)
97-
98-
assert mock_context.set_checkpoint.call_count == 3
99-
100-
calls = mock_context.set_checkpoint.call_args_list
101-
expected_arns = [
102-
"arn:aws:sqs:us-east-1:123456789012:queue1",
103-
"arn:aws:sqs:us-east-1:123456789012:queue2",
104-
"arn:aws:sqs:us-east-1:123456789012:queue3",
105-
]
106-
107-
for i, call in enumerate(calls):
108-
args, kwargs = call
109-
tags = args[0]
110-
self.assertIn("direction:in", tags)
111-
self.assertIn(f"topic:{expected_arns[i]}", tags)
112-
self.assertIn("type:sqs", tags)
113-
self.assertEqual(kwargs["payload_size"], 100)
110+
for i, call in enumerate(calls):
111+
args, kwargs = call
112+
tags = args[0]
113+
self.assertIn("direction:in", tags)
114+
self.assertIn(f"topic:{expected_arns[i]}", tags)
115+
self.assertIn("type:sqs", tags)
116+
self.assertEqual(kwargs["payload_size"], 100)

‎tests/test_wrapper.py‎

Lines changed: 46 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,10 @@ def setUp(self):
7676
self.mock_dd_lambda_layer_tag = patcher.start()
7777
self.addCleanup(patcher.stop)
7878

79+
patcher = patch("datadog_lambda.wrapper.set_dsm_context")
80+
self.mock_set_dsm_context = patcher.start()
81+
self.addCleanup(patcher.stop)
82+
7983
def test_datadog_lambda_wrapper(self):
8084
wrapper.dd_tracing_enabled = False
8185

@@ -564,68 +568,60 @@ def return_type_test(event, context):
564568
self.assertFalse(MockPrintExc.called)
565569

566570
def test_set_dsm_context_called_when_DSM_and_tracing_enabled(self):
567-
env_vars = {"DD_DATA_STREAMS_ENABLED": "true"}
568-
with patch.dict(os.environ, env_vars):
569-
with patch("datadog_lambda.wrapper.dd_tracing_enabled", True):
570-
with patch(
571-
"datadog_lambda.wrapper.set_dsm_context"
572-
) as set_dsm_context_patch:
571+
os.environ["DD_DATA_STREAMS_ENABLED"] = "true"
572+
wrapper.dd_tracing_enabled = True
573+
574+
@wrapper.datadog_lambda_wrapper
575+
def lambda_handler(event, context):
576+
return "ok"
573577

574-
@wrapper.datadog_lambda_wrapper
575-
def lambda_handler(event, context):
576-
return "ok"
578+
result = lambda_handler({}, get_mock_context())
579+
self.assertEqual(result, "ok")
580+
self.mock_set_dsm_context.assert_called_once()
577581

578-
result = lambda_handler({}, get_mock_context())
579-
assert result == "ok"
580-
assert set_dsm_context_patch.call_count == 1
582+
del os.environ["DD_DATA_STREAMS_ENABLED"]
581583

582584
def test_set_dsm_context_not_called_when_only_DSM_enabled(self):
583-
env_vars = {"DD_DATA_STREAMS_ENABLED": "true"}
584-
with patch.dict(os.environ, env_vars):
585-
with patch("datadog_lambda.wrapper.dd_tracing_enabled", False):
586-
with patch(
587-
"datadog_lambda.wrapper.set_dsm_context"
588-
) as set_dsm_context_patch:
585+
os.environ["DD_DATA_STREAMS_ENABLED"] = "true"
586+
wrapper.dd_tracing_enabled = False
587+
588+
@wrapper.datadog_lambda_wrapper
589+
def lambda_handler(event, context):
590+
return "ok"
589591

590-
@wrapper.datadog_lambda_wrapper
591-
def lambda_handler(event, context):
592-
return "ok"
592+
result = lambda_handler({}, get_mock_context())
593+
self.assertEqual(result, "ok")
594+
self.mock_set_dsm_context.assert_not_called()
593595

594-
result = lambda_handler({}, get_mock_context())
595-
assert result == "ok"
596-
assert set_dsm_context_patch.call_count == 0
596+
del os.environ["DD_DATA_STREAMS_ENABLED"]
597597

598598
def test_set_dsm_context_not_called_when_only_tracing_enabled(self):
599-
env_vars = {"DD_DATA_STREAMS_ENABLED": "false"}
600-
with patch.dict(os.environ, env_vars):
601-
with patch("datadog_lambda.wrapper.dd_tracing_enabled", True):
602-
with patch(
603-
"datadog_lambda.wrapper.set_dsm_context"
604-
) as set_dsm_context_patch:
599+
os.environ["DD_DATA_STREAMS_ENABLED"] = "false"
600+
wrapper.dd_tracing_enabled = True
605601

606-
@wrapper.datadog_lambda_wrapper
607-
def lambda_handler(event, context):
608-
return "ok"
602+
@wrapper.datadog_lambda_wrapper
603+
def lambda_handler(event, context):
604+
return "ok"
609605

610-
result = lambda_handler({}, get_mock_context())
611-
assert result == "ok"
612-
assert set_dsm_context_patch.call_count == 0
606+
result = lambda_handler({}, get_mock_context())
607+
self.assertEqual(result, "ok")
608+
self.mock_set_dsm_context.assert_not_called()
609+
610+
del os.environ["DD_DATA_STREAMS_ENABLED"]
613611

614612
def test_set_dsm_context_not_called_when_tracing_and_DSM_disabled(self):
615-
env_vars = {"DD_DATA_STREAMS_ENABLED": "false"}
616-
with patch.dict(os.environ, env_vars):
617-
with patch("datadog_lambda.wrapper.dd_tracing_enabled", False):
618-
with patch(
619-
"datadog_lambda.wrapper.set_dsm_context"
620-
) as set_dsm_context_patch:
621-
622-
@wrapper.datadog_lambda_wrapper
623-
def lambda_handler(event, context):
624-
return "ok"
625-
626-
result = lambda_handler({}, get_mock_context())
627-
assert result == "ok"
628-
assert set_dsm_context_patch.call_count == 0
613+
os.environ["DD_DATA_STREAMS_ENABLED"] = "false"
614+
wrapper.dd_tracing_enabled = False
615+
616+
@wrapper.datadog_lambda_wrapper
617+
def lambda_handler(event, context):
618+
return "ok"
619+
620+
result = lambda_handler({}, get_mock_context())
621+
self.assertEqual(result, "ok")
622+
self.mock_set_dsm_context.assert_not_called()
623+
624+
del os.environ["DD_DATA_STREAMS_ENABLED"]
629625

630626

631627
class TestLambdaDecoratorSettings(unittest.TestCase):

0 commit comments

Comments
 (0)