7
7
from datadog .api .exceptions import ClientError
8
8
from datetime import datetime , timedelta
9
9
10
- from datadog_lambda .metric import lambda_metric
10
+ from datadog_lambda .metric import lambda_metric , flush_stats
11
11
from datadog_lambda .api import decrypt_kms_api_key , KMS_ENCRYPTION_CONTEXT_KEY
12
12
from datadog_lambda .thread_stats_writer import ThreadStatsWriter
13
13
from datadog_lambda .tags import dd_lambda_layer_tag
@@ -101,6 +101,10 @@ def setUp(self):
101
101
self .mock_threadstats_flush_distributions = patcher .start ()
102
102
self .addCleanup (patcher .stop )
103
103
104
+ patcher = patch ("datadog_lambda.metric.extension_thread_stats" )
105
+ self .mock_extension_thread_stats = patcher .start ()
106
+ self .addCleanup (patcher .stop )
107
+
104
108
def test_retry_on_remote_disconnected (self ):
105
109
# Raise the RemoteDisconnected error
106
110
lambda_stats = ThreadStatsWriter (True )
@@ -123,6 +127,10 @@ def test_flush_stats_with_tags(self):
123
127
for tag in tags :
124
128
self .assertTrue (tag in lambda_stats .thread_stats .constant_tags )
125
129
130
+ def test_flush_stats_without_context (self ):
131
+ flush_stats (lambda_context = None )
132
+ self .mock_extension_thread_stats .flush .assert_called_with (None )
133
+
126
134
127
135
MOCK_FUNCTION_NAME = "myFunction"
128
136
0 commit comments