Skip to content

Commit 6f65160

Browse files
michael-zhao459purple4reina
authored andcommitted
fixes
1 parent dda744e commit 6f65160

File tree

1 file changed

+11
-15
lines changed

1 file changed

+11
-15
lines changed

tests/test_dsm.py

Lines changed: 11 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,10 @@
22
from unittest.mock import patch, MagicMock
33

44
from datadog_lambda.dsm import set_dsm_context, _dsm_set_sqs_context
5-
from datadog_lambda.trigger import EventTypes
5+
from datadog_lambda.trigger import EventTypes, _EventSource
66

77

8-
class TestDsmContext(unittest.TestCase):
8+
class TestDsmSQSContext(unittest.TestCase):
99
def setUp(self):
1010
patcher = patch("datadog_lambda.dsm._dsm_set_sqs_context")
1111
self.mock_dsm_set_sqs_context = patcher.start()
@@ -33,17 +33,15 @@ def setUp(self):
3333

3434
def test_non_sqs_event_source_does_nothing(self):
3535
"""Test that non-SQS event sources don't trigger DSM context setting"""
36-
event = {"Records": [{"body": "test"}]}
36+
event = {}
37+
# Use Unknown Event Source
38+
event_source = _EventSource(EventTypes.UNKNOWN)
39+
set_dsm_context(event, event_source)
3740

38-
mock_event_source = MagicMock()
39-
mock_event_source.equals.return_value = False # Not SQS
40-
41-
set_dsm_context(event, mock_event_source)
42-
43-
mock_event_source.equals.assert_called_once_with(EventTypes.SQS)
41+
# DSM context should not be set for non-SQS events
4442
self.mock_dsm_set_sqs_context.assert_not_called()
4543

46-
def test_event_with_no_records_does_nothing(self):
44+
def test_sqs_event_with_no_records_does_nothing(self):
4745
"""Test that events where Records is None don't trigger DSM processing"""
4846
events_with_no_records = [
4947
{},
@@ -67,14 +65,12 @@ def test_sqs_event_triggers_dsm_sqs_context(self):
6765
]
6866
}
6967

70-
mock_event_source = MagicMock()
71-
mock_event_source.equals.return_value = True
72-
73-
set_dsm_context(sqs_event, mock_event_source)
68+
event_source = _EventSource(EventTypes.SQS)
69+
set_dsm_context(sqs_event, event_source)
7470

7571
self.mock_dsm_set_sqs_context.assert_called_once_with(sqs_event)
7672

77-
def test_multiple_records_process_each_record(self):
73+
def test_sqs_multiple_records_process_each_record(self):
7874
"""Test that each record in an SQS event gets processed individually"""
7975
multi_record_event = {
8076
"Records": [

0 commit comments

Comments
 (0)