77from datadog .api .exceptions import ClientError
88from datetime import datetime , timedelta
99
10- from datadog_lambda .metric import lambda_metric
10+ from datadog_lambda .metric import lambda_metric , flush_stats
1111from datadog_lambda .api import decrypt_kms_api_key , KMS_ENCRYPTION_CONTEXT_KEY
1212from datadog_lambda .thread_stats_writer import ThreadStatsWriter
1313from datadog_lambda .tags import dd_lambda_layer_tag
@@ -101,6 +101,10 @@ def setUp(self):
101101 self .mock_threadstats_flush_distributions = patcher .start ()
102102 self .addCleanup (patcher .stop )
103103
104+ patcher = patch ("datadog_lambda.metric.extension_thread_stats" )
105+ self .mock_extension_thread_stats = patcher .start ()
106+ self .addCleanup (patcher .stop )
107+
104108 def test_retry_on_remote_disconnected (self ):
105109 # Raise the RemoteDisconnected error
106110 lambda_stats = ThreadStatsWriter (True )
@@ -123,6 +127,10 @@ def test_flush_stats_with_tags(self):
123127 for tag in tags :
124128 self .assertTrue (tag in lambda_stats .thread_stats .constant_tags )
125129
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+
126134
127135MOCK_FUNCTION_NAME = "myFunction"
128136
0 commit comments