@@ -38,14 +38,6 @@ def setUp(self):
38
38
self .mock_threadstats_flush_distributions = patcher .start ()
39
39
self .addCleanup (patcher .stop )
40
40
41
- self .metric_lambda_stats_patcher = patch ("datadog_lambda.metric.lambda_stats" )
42
- self .mock_metric_lambda_stats = self .metric_lambda_stats_patcher .start ()
43
- self .addCleanup (self .metric_lambda_stats_patcher .stop )
44
-
45
- self .wrapper_lambda_stats_patcher = patch ("datadog_lambda.wrapper.lambda_stats" )
46
- self .mock_wrapper_lambda_stats = self .wrapper_lambda_stats_patcher .start ()
47
- self .addCleanup (self .wrapper_lambda_stats_patcher .stop )
48
-
49
41
patcher = patch ("datadog_lambda.wrapper.extract_dd_trace_context" )
50
42
self .mock_extract_dd_trace_context = patcher .start ()
51
43
self .mock_extract_dd_trace_context .return_value = ({}, None )
@@ -101,10 +93,16 @@ def lambda_handler(event, context):
101
93
102
94
lambda_handler (lambda_event , lambda_context )
103
95
104
- self .mock_metric_lambda_stats .distribution .assert_has_calls (
105
- [call ("test.metric" , 100 , timestamp = None , tags = ANY )]
96
+ self .mock_threadstats_flush_distributions .assert_has_calls (
97
+ [call ([{
98
+ 'metric' : 'test.metric' , 'points' : [[ANY , [100 ]]],
99
+ 'type' : 'distribution' ,
100
+ 'host' : None ,
101
+ 'device' : None ,
102
+ 'tags' : ANY ,
103
+ 'interval' : 10
104
+ }])]
106
105
)
107
- self .mock_wrapper_lambda_stats .flush .assert_called ()
108
106
self .mock_extract_dd_trace_context .assert_called_with (
109
107
lambda_event , lambda_context , extractor = None
110
108
)
@@ -122,16 +120,11 @@ def lambda_handler(event, context):
122
120
lambda_event = {}
123
121
lambda_handler (lambda_event , get_mock_context ())
124
122
125
- self .mock_metric_lambda_stats .distribution .assert_not_called ()
126
- self .mock_wrapper_lambda_stats .flush .assert_not_called ()
123
+ self .mock_threadstats_flush_distributions .assert_not_called ()
127
124
128
125
del os .environ ["DD_FLUSH_TO_LOG" ]
129
126
130
127
def test_datadog_lambda_wrapper_flush_in_thread (self ):
131
- # stop patchers so mock_threadstats_flush_distributions gets called
132
- self .metric_lambda_stats_patcher .stop ()
133
- self .wrapper_lambda_stats_patcher .stop ()
134
-
135
128
# force ThreadStats to flush in thread
136
129
import datadog_lambda .metric as metric_module
137
130
@@ -158,15 +151,7 @@ def lambda_handler(event, context):
158
151
metric_module .lambda_stats .stop ()
159
152
metric_module .lambda_stats .start (flush_in_thread = False )
160
153
161
- # resume patchers
162
- self .metric_lambda_stats_patcher .start ()
163
- self .wrapper_lambda_stats_patcher .start ()
164
-
165
154
def test_datadog_lambda_wrapper_not_flush_in_thread (self ):
166
- # stop patchers so mock_threadstats_flush_distributions gets called
167
- self .metric_lambda_stats_patcher .stop ()
168
- self .wrapper_lambda_stats_patcher .stop ()
169
-
170
155
# force ThreadStats to not flush in thread
171
156
import datadog_lambda .metric as metric_module
172
157
@@ -193,10 +178,6 @@ def lambda_handler(event, context):
193
178
metric_module .lambda_stats .stop ()
194
179
metric_module .lambda_stats .start (flush_in_thread = False )
195
180
196
- # resume patchers
197
- self .metric_lambda_stats_patcher .start ()
198
- self .wrapper_lambda_stats_patcher .start ()
199
-
200
181
def test_datadog_lambda_wrapper_inject_correlation_ids (self ):
201
182
os .environ ["DD_LOGS_INJECTION" ] = "True"
202
183
@@ -445,5 +426,4 @@ def lambda_handler(event, context):
445
426
lambda_handler_double_wrapped (lambda_event , get_mock_context ())
446
427
447
428
self .mock_patch_all .assert_called_once ()
448
- self .mock_wrapper_lambda_stats .flush .assert_called_once ()
449
429
self .mock_submit_invocations_metric .assert_called_once ()
0 commit comments