diff --git a/datadog_lambda/metric.py b/datadog_lambda/metric.py index 44f8e377..d312e3bb 100644 --- a/datadog_lambda/metric.py +++ b/datadog_lambda/metric.py @@ -125,6 +125,7 @@ def flush_stats(lambda_context=None): lambda_stats.flush() if extension_thread_stats is not None: + tags = None if lambda_context is not None: tags = get_enhanced_metrics_tags(lambda_context) split_arn = lambda_context.invoked_function_arn.split(":") diff --git a/tests/test_metric.py b/tests/test_metric.py index 10ef9191..241f563b 100644 --- a/tests/test_metric.py +++ b/tests/test_metric.py @@ -7,7 +7,7 @@ from datadog.api.exceptions import ClientError from datetime import datetime, timedelta -from datadog_lambda.metric import lambda_metric +from datadog_lambda.metric import lambda_metric, flush_stats from datadog_lambda.api import decrypt_kms_api_key, KMS_ENCRYPTION_CONTEXT_KEY from datadog_lambda.thread_stats_writer import ThreadStatsWriter from datadog_lambda.tags import dd_lambda_layer_tag @@ -101,6 +101,10 @@ def setUp(self): self.mock_threadstats_flush_distributions = patcher.start() self.addCleanup(patcher.stop) + patcher = patch("datadog_lambda.metric.extension_thread_stats") + self.mock_extension_thread_stats = patcher.start() + self.addCleanup(patcher.stop) + def test_retry_on_remote_disconnected(self): # Raise the RemoteDisconnected error lambda_stats = ThreadStatsWriter(True) @@ -123,6 +127,10 @@ def test_flush_stats_with_tags(self): for tag in tags: self.assertTrue(tag in lambda_stats.thread_stats.constant_tags) + def test_flush_stats_without_context(self): + flush_stats(lambda_context=None) + self.mock_extension_thread_stats.flush.assert_called_with(None) + MOCK_FUNCTION_NAME = "myFunction"