Skip to content

Commit de6b4ff

Browse files
authored
Support calling flush_stats with None context. (#518)
* Support calling flush_stats with None context. * Fix lint error.
1 parent adfc9e8 commit de6b4ff

File tree

2 files changed

+10
-1
lines changed

2 files changed

+10
-1
lines changed

datadog_lambda/metric.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,7 @@ def flush_stats(lambda_context=None):
125125
lambda_stats.flush()
126126

127127
if extension_thread_stats is not None:
128+
tags = None
128129
if lambda_context is not None:
129130
tags = get_enhanced_metrics_tags(lambda_context)
130131
split_arn = lambda_context.invoked_function_arn.split(":")

tests/test_metric.py

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
from datadog.api.exceptions import ClientError
88
from datetime import datetime, timedelta
99

10-
from datadog_lambda.metric import lambda_metric
10+
from datadog_lambda.metric import lambda_metric, flush_stats
1111
from datadog_lambda.api import decrypt_kms_api_key, KMS_ENCRYPTION_CONTEXT_KEY
1212
from datadog_lambda.thread_stats_writer import ThreadStatsWriter
1313
from 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

127135
MOCK_FUNCTION_NAME = "myFunction"
128136

0 commit comments

Comments
 (0)