From a822cd83863ced000365f5eed66eeb8781138609 Mon Sep 17 00:00:00 2001 From: Rey Abolofia Date: Fri, 27 Sep 2024 14:08:27 -0700 Subject: [PATCH 1/2] Support calling flush_stats with None context. --- datadog_lambda/metric.py | 1 + tests/test_metric.py | 12 +++++++++++- 2 files changed, 12 insertions(+), 1 deletion(-) 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..a76a1022 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,12 @@ 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 +129,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" From 742038c86769fabb1bc34b117f594ad405416425 Mon Sep 17 00:00:00 2001 From: Rey Abolofia Date: Mon, 30 Sep 2024 10:04:45 -0700 Subject: [PATCH 2/2] Fix lint error. --- tests/test_metric.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/test_metric.py b/tests/test_metric.py index a76a1022..241f563b 100644 --- a/tests/test_metric.py +++ b/tests/test_metric.py @@ -101,9 +101,7 @@ def setUp(self): self.mock_threadstats_flush_distributions = patcher.start() self.addCleanup(patcher.stop) - patcher = patch( - "datadog_lambda.metric.extension_thread_stats" - ) + patcher = patch("datadog_lambda.metric.extension_thread_stats") self.mock_extension_thread_stats = patcher.start() self.addCleanup(patcher.stop)