2
2
from unittest .mock import patch , MagicMock
3
3
4
4
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
6
6
7
7
8
- class TestDsmContext (unittest .TestCase ):
8
+ class TestDsmSQSContext (unittest .TestCase ):
9
9
def setUp (self ):
10
10
patcher = patch ("datadog_lambda.dsm._dsm_set_sqs_context" )
11
11
self .mock_dsm_set_sqs_context = patcher .start ()
@@ -33,17 +33,15 @@ def setUp(self):
33
33
34
34
def test_non_sqs_event_source_does_nothing (self ):
35
35
"""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 )
37
40
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
44
42
self .mock_dsm_set_sqs_context .assert_not_called ()
45
43
46
- def test_event_with_no_records_does_nothing (self ):
44
+ def test_sqs_event_with_no_records_does_nothing (self ):
47
45
"""Test that events where Records is None don't trigger DSM processing"""
48
46
events_with_no_records = [
49
47
{},
@@ -67,14 +65,12 @@ def test_sqs_event_triggers_dsm_sqs_context(self):
67
65
]
68
66
}
69
67
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 )
74
70
75
71
self .mock_dsm_set_sqs_context .assert_called_once_with (sqs_event )
76
72
77
- def test_multiple_records_process_each_record (self ):
73
+ def test_sqs_multiple_records_process_each_record (self ):
78
74
"""Test that each record in an SQS event gets processed individually"""
79
75
multi_record_event = {
80
76
"Records" : [
0 commit comments