From 6c58173e543c2d4ac032a30b8f43694af41b3f95 Mon Sep 17 00:00:00 2001 From: Aleksandr Pasechnik Date: Wed, 30 Apr 2025 15:40:41 -0400 Subject: [PATCH 01/10] feat: dogstatsd client sends timestamp --- datadog_lambda/dogstatsd.py | 23 +++++++++++++---------- tests/test_dogstatsd.py | 20 ++++++++++++-------- 2 files changed, 25 insertions(+), 18 deletions(-) diff --git a/datadog_lambda/dogstatsd.py b/datadog_lambda/dogstatsd.py index a627492d..b9c0b697 100644 --- a/datadog_lambda/dogstatsd.py +++ b/datadog_lambda/dogstatsd.py @@ -1,11 +1,10 @@ +import errno import logging import os -import socket -import errno import re +import socket from threading import Lock - MIN_SEND_BUFFER_SIZE = 32 * 1024 log = logging.getLogger("datadog_lambda.dogstatsd") @@ -55,14 +54,17 @@ def _get_udp_socket(cls, host, port): return sock - def distribution(self, metric, value, tags=None): + def distribution(self, metric, value, tags=None, timestamp=None): """ - Send a global distribution value, optionally setting tags. + Send a global distribution value, optionally setting tags. The optional + timestamp should be an integer representing seconds since the epoch + (January 1, 1970, 00:00:00 UTC). >>> statsd.distribution("uploaded.file.size", 1445) >>> statsd.distribution("album.photo.count", 26, tags=["gender:female"]) + >>> statsd.distribution("historic.file.count", 5, timestamp=int(datetime(2020, 2, 14, 12, 0, 0).timestamp())) """ - self._report(metric, "d", value, tags) + self._report(metric, "d", value, tags, timestamp) def close_socket(self): """ @@ -84,20 +86,21 @@ def normalize_tags(self, tag_list): for tag in tag_list ] - def _serialize_metric(self, metric, metric_type, value, tags): + def _serialize_metric(self, metric, metric_type, value, tags, timestamp): # Create/format the metric packet - return "%s:%s|%s%s" % ( + return "%s:%s|%s%s%s" % ( metric, value, metric_type, ("|#" + ",".join(self.normalize_tags(tags))) if tags else "", + ("|T" + str(timestamp)) if timestamp is not None else "", ) - def _report(self, metric, metric_type, value, tags): + def _report(self, metric, metric_type, value, tags, timestamp): if value is None: return - payload = self._serialize_metric(metric, metric_type, value, tags) + payload = self._serialize_metric(metric, metric_type, value, tags, timestamp) # Send it self._send_to_server(payload) diff --git a/tests/test_dogstatsd.py b/tests/test_dogstatsd.py index 149e1a70..ea6afd48 100644 --- a/tests/test_dogstatsd.py +++ b/tests/test_dogstatsd.py @@ -1,5 +1,5 @@ -from collections import deque import unittest +from collections import deque from datadog_lambda.dogstatsd import statsd @@ -36,16 +36,20 @@ def test_init(self): self.assertEqual(statsd.port, 8125) self.assertEqual(statsd.encoding, "utf-8") - def test_distribution_no_tags(self): - statsd.distribution("my.test.metric", 3) + def _checkOnlyOneMetric(self, value): payload = self.recv() metrics = payload.split("\n") self.assertEqual(len(metrics), 1) - self.assertEqual("my.test.metric:3|d", metrics[0]) + self.assertEqual(value, metrics[0]) + + def test_distribution_no_tags(self): + statsd.distribution("my.test.metric", 3) + self._checkOnlyOneMetric("my.test.metric:3|d") def test_distribution_with_tags(self): statsd.distribution("my.test.tags.metric", 3, tags=["taga:valuea,tagb:valueb"]) - payload = self.recv() - metrics = payload.split("\n") - self.assertEqual(len(metrics), 1) - self.assertEqual("my.test.tags.metric:3|d|#taga:valuea_tagb:valueb", metrics[0]) + self._checkOnlyOneMetric("my.test.tags.metric:3|d|#taga:valuea_tagb:valueb") + + def test_distribution_with_timestamp(self): + statsd.distribution("my.test.timestamp.metric", 9, timestamp=123456789) + self._checkOnlyOneMetric("my.test.timestamp.metric:9|d|T123456789") From b4f6c81dd1972807844cf46a598688dc167dbef7 Mon Sep 17 00:00:00 2001 From: Aleksandr Pasechnik Date: Wed, 30 Apr 2025 15:57:36 -0400 Subject: [PATCH 02/10] chore: wire up timestamp to StatsDWriter and fix tags=[] bug --- datadog_lambda/metric.py | 15 +++++++++------ datadog_lambda/stats_writer.py | 2 +- datadog_lambda/statsd_writer.py | 6 +++--- datadog_lambda/thread_stats_writer.py | 3 ++- 4 files changed, 15 insertions(+), 11 deletions(-) diff --git a/datadog_lambda/metric.py b/datadog_lambda/metric.py index f9c67a26..5bf275e5 100644 --- a/datadog_lambda/metric.py +++ b/datadog_lambda/metric.py @@ -3,14 +3,15 @@ # This product includes software developed at Datadog (https://www.datadoghq.com/). # Copyright 2019 Datadog, Inc. +import logging import os import time -import logging -import ujson as json from datetime import datetime, timedelta +import ujson as json + from datadog_lambda.extension import should_use_extension -from datadog_lambda.tags import get_enhanced_metrics_tags, dd_lambda_layer_tag +from datadog_lambda.tags import dd_lambda_layer_tag, get_enhanced_metrics_tags logger = logging.getLogger(__name__) @@ -28,8 +29,8 @@ # and leads to data loss. When disabled, metrics are only flushed at the # end of invocation. To make metrics submitted from a long-running Lambda # function available sooner, consider using the Datadog Lambda extension. - from datadog_lambda.thread_stats_writer import ThreadStatsWriter from datadog_lambda.api import init_api + from datadog_lambda.thread_stats_writer import ThreadStatsWriter init_api() lambda_stats = ThreadStatsWriter(flush_in_thread) @@ -92,8 +93,8 @@ def lambda_metric(metric_name, value, timestamp=None, tags=None, force_async=Fal return global extension_thread_stats if extension_thread_stats is None: - from datadog_lambda.thread_stats_writer import ThreadStatsWriter from datadog_lambda.api import init_api + from datadog_lambda.thread_stats_writer import ThreadStatsWriter init_api() extension_thread_stats = ThreadStatsWriter(flush_in_thread) @@ -119,8 +120,10 @@ def lambda_metric(metric_name, value, timestamp=None, tags=None, force_async=Fal ) -def write_metric_point_to_stdout(metric_name, value, timestamp=None, tags=[]): +def write_metric_point_to_stdout(metric_name, value, timestamp=None, tags=None): """Writes the specified metric point to standard output""" + tags = tags or [] + logger.debug( "Sending metric %s value %s to Datadog via log forwarder", metric_name, value ) diff --git a/datadog_lambda/stats_writer.py b/datadog_lambda/stats_writer.py index d3919c30..563b1ae9 100644 --- a/datadog_lambda/stats_writer.py +++ b/datadog_lambda/stats_writer.py @@ -1,5 +1,5 @@ class StatsWriter: - def distribution(self, metric_name, value, tags=[], timestamp=None): + def distribution(self, metric_name, value, tags=None, timestamp=None): raise NotImplementedError() def flush(self): diff --git a/datadog_lambda/statsd_writer.py b/datadog_lambda/statsd_writer.py index 33843dc6..4aaab8d5 100644 --- a/datadog_lambda/statsd_writer.py +++ b/datadog_lambda/statsd_writer.py @@ -1,5 +1,5 @@ -from datadog_lambda.stats_writer import StatsWriter from datadog_lambda.dogstatsd import statsd +from datadog_lambda.stats_writer import StatsWriter class StatsDWriter(StatsWriter): @@ -7,8 +7,8 @@ class StatsDWriter(StatsWriter): Writes distribution metrics using StatsD protocol """ - def distribution(self, metric_name, value, tags=[], timestamp=None): - statsd.distribution(metric_name, value, tags=tags) + def distribution(self, metric_name, value, tags=None, timestamp=None): + statsd.distribution(metric_name, value, tags=tags, timestamp=timestamp) def flush(self): pass diff --git a/datadog_lambda/thread_stats_writer.py b/datadog_lambda/thread_stats_writer.py index 422a9a0a..f21ee31f 100644 --- a/datadog_lambda/thread_stats_writer.py +++ b/datadog_lambda/thread_stats_writer.py @@ -3,6 +3,7 @@ # Make sure that this package would always be lazy-loaded/outside from the critical path # since underlying packages are quite heavy to load and useless when the extension is present from datadog.threadstats import ThreadStats + from datadog_lambda.stats_writer import StatsWriter logger = logging.getLogger(__name__) @@ -17,7 +18,7 @@ def __init__(self, flush_in_thread): self.thread_stats = ThreadStats(compress_payload=True) self.thread_stats.start(flush_in_thread=flush_in_thread) - def distribution(self, metric_name, value, tags=[], timestamp=None): + def distribution(self, metric_name, value, tags=None, timestamp=None): self.thread_stats.distribution( metric_name, value, tags=tags, timestamp=timestamp ) From 1e564c967568b6754f9cb7ed8c6a611adb5a34ac Mon Sep 17 00:00:00 2001 From: Aleksandr Pasechnik Date: Fri, 2 May 2025 13:20:58 -0400 Subject: [PATCH 03/10] feat: send metrics with timestamps to the extension This is now supported by both the go extension and bottlecap, so we can safely send metrics to the extension and not have to ship the ones with timestamps directly to the Datadog API. --- datadog_lambda/metric.py | 42 ++++++++++++++-------------------------- tests/test_metric.py | 32 ++++++++++++++++++------------ 2 files changed, 34 insertions(+), 40 deletions(-) diff --git a/datadog_lambda/metric.py b/datadog_lambda/metric.py index 5bf275e5..7817a5a0 100644 --- a/datadog_lambda/metric.py +++ b/datadog_lambda/metric.py @@ -76,39 +76,25 @@ def lambda_metric(metric_name, value, timestamp=None, tags=None, force_async=Fal tags = [] if tags is None else list(tags) tags.append(dd_lambda_layer_tag) - if should_use_extension and timestamp is not None: - # The extension does not support timestamps for distributions so we create a - # a thread stats writer to submit metrics with timestamps to the API - timestamp_ceiling = int( - (datetime.now() - timedelta(hours=4)).timestamp() - ) # 4 hours ago - if isinstance(timestamp, datetime): - timestamp = int(timestamp.timestamp()) - if timestamp_ceiling > timestamp: - logger.warning( - "Timestamp %s is older than 4 hours, not submitting metric %s", - timestamp, - metric_name, - ) - return - global extension_thread_stats - if extension_thread_stats is None: - from datadog_lambda.api import init_api - from datadog_lambda.thread_stats_writer import ThreadStatsWriter - - init_api() - extension_thread_stats = ThreadStatsWriter(flush_in_thread) - - extension_thread_stats.distribution( - metric_name, value, tags=tags, timestamp=timestamp - ) - return - if should_use_extension: + if timestamp is not None: + if isinstance(timestamp, datetime): + timestamp = int(timestamp.timestamp()) + + timestamp_floor = int((datetime.now() - timedelta(hours=4)).timestamp()) + if timestamp < timestamp_floor: + logger.warning( + "Timestamp %s is older than 4 hours, not submitting metric %s", + timestamp, + metric_name, + ) + return + logger.debug( "Sending metric %s value %s to Datadog via extension", metric_name, value ) lambda_stats.distribution(metric_name, value, tags=tags, timestamp=timestamp) + else: if flush_to_logs or force_async: write_metric_point_to_stdout( diff --git a/tests/test_metric.py b/tests/test_metric.py index d10a0f0d..4caf6035 100644 --- a/tests/test_metric.py +++ b/tests/test_metric.py @@ -1,16 +1,15 @@ import os import unittest - -from unittest.mock import patch, call +from datetime import datetime, timedelta +from unittest.mock import call, patch from botocore.exceptions import ClientError as BotocoreClientError from datadog.api.exceptions import ClientError -from datetime import datetime, timedelta -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.api import KMS_ENCRYPTION_CONTEXT_KEY, decrypt_kms_api_key +from datadog_lambda.metric import flush_stats, lambda_metric from datadog_lambda.tags import dd_lambda_layer_tag +from datadog_lambda.thread_stats_writer import ThreadStatsWriter class TestLambdaMetric(unittest.TestCase): @@ -53,10 +52,10 @@ def test_lambda_metric_timestamp_with_extension(self): timestamp = int((datetime.now() - delta).timestamp()) lambda_metric("test_timestamp", 1, timestamp) - self.mock_metric_lambda_stats.distribution.assert_not_called() - self.mock_metric_extension_thread_stats.distribution.assert_called_with( - "test_timestamp", 1, timestamp=timestamp, tags=[dd_lambda_layer_tag] + self.mock_metric_lambda_stats.distribution.assert_has_calls( + [call("test_timestamp", 1, timestamp=timestamp, tags=[dd_lambda_layer_tag])] ) + self.mock_metric_extension_thread_stats.distribution.assert_not_called() @patch("datadog_lambda.metric.should_use_extension", True) def test_lambda_metric_datetime_with_extension(self): @@ -64,11 +63,20 @@ def test_lambda_metric_datetime_with_extension(self): self.mock_metric_extension_thread_stats = patcher.start() self.addCleanup(patcher.stop) - delta = timedelta(hours=5) + delta = timedelta(minutes=1) timestamp = datetime.now() - delta - lambda_metric("test_timestamp", 1, timestamp) - self.mock_metric_lambda_stats.distribution.assert_not_called() + lambda_metric("test_datetime_timestamp", 0, timestamp) + self.mock_metric_lambda_stats.distribution.assert_has_calls( + [ + call( + "test_datetime_timestamp", + 0, + timestamp=int(timestamp.timestamp()), + tags=[dd_lambda_layer_tag], + ) + ] + ) self.mock_metric_extension_thread_stats.distribution.assert_not_called() @patch("datadog_lambda.metric.should_use_extension", True) From ee1249ae96f175f8384c6b2e9fa14862794975cc Mon Sep 17 00:00:00 2001 From: Aleksandr Pasechnik Date: Fri, 2 May 2025 13:36:56 -0400 Subject: [PATCH 04/10] chore: remove unnecessary extension_thread_stats --- datadog_lambda/dogstatsd.py | 6 +++++- datadog_lambda/metric.py | 13 ------------- tests/test_metric.py | 23 ----------------------- 3 files changed, 5 insertions(+), 37 deletions(-) diff --git a/datadog_lambda/dogstatsd.py b/datadog_lambda/dogstatsd.py index b9c0b697..f30a2039 100644 --- a/datadog_lambda/dogstatsd.py +++ b/datadog_lambda/dogstatsd.py @@ -62,7 +62,11 @@ def distribution(self, metric, value, tags=None, timestamp=None): >>> statsd.distribution("uploaded.file.size", 1445) >>> statsd.distribution("album.photo.count", 26, tags=["gender:female"]) - >>> statsd.distribution("historic.file.count", 5, timestamp=int(datetime(2020, 2, 14, 12, 0, 0).timestamp())) + >>> statsd.distribution( + >>> "historic.file.count", + >>> 5, + >>> timestamp=int(datetime(2020, 2, 14, 12, 0, 0).timestamp()), + >>> ) """ self._report(metric, "d", value, tags, timestamp) diff --git a/datadog_lambda/metric.py b/datadog_lambda/metric.py index 7817a5a0..79e512f8 100644 --- a/datadog_lambda/metric.py +++ b/datadog_lambda/metric.py @@ -16,7 +16,6 @@ logger = logging.getLogger(__name__) lambda_stats = None -extension_thread_stats = None flush_in_thread = os.environ.get("DD_FLUSH_IN_THREAD", "").lower() == "true" @@ -129,18 +128,6 @@ def write_metric_point_to_stdout(metric_name, value, timestamp=None, tags=None): 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(":") - if len(split_arn) > 7: - # Get rid of the alias - split_arn.pop() - arn = ":".join(split_arn) - tags.append("function_arn:" + arn) - extension_thread_stats.flush(tags) - def submit_enhanced_metric(metric_name, lambda_context): """Submits the enhanced metric with the given name diff --git a/tests/test_metric.py b/tests/test_metric.py index 4caf6035..d32c0876 100644 --- a/tests/test_metric.py +++ b/tests/test_metric.py @@ -44,10 +44,6 @@ def test_lambda_metric_flush_to_log_with_extension(self): @patch("datadog_lambda.metric.should_use_extension", True) def test_lambda_metric_timestamp_with_extension(self): - patcher = patch("datadog_lambda.metric.extension_thread_stats") - self.mock_metric_extension_thread_stats = patcher.start() - self.addCleanup(patcher.stop) - delta = timedelta(minutes=1) timestamp = int((datetime.now() - delta).timestamp()) @@ -55,14 +51,9 @@ def test_lambda_metric_timestamp_with_extension(self): self.mock_metric_lambda_stats.distribution.assert_has_calls( [call("test_timestamp", 1, timestamp=timestamp, tags=[dd_lambda_layer_tag])] ) - self.mock_metric_extension_thread_stats.distribution.assert_not_called() @patch("datadog_lambda.metric.should_use_extension", True) def test_lambda_metric_datetime_with_extension(self): - patcher = patch("datadog_lambda.metric.extension_thread_stats") - self.mock_metric_extension_thread_stats = patcher.start() - self.addCleanup(patcher.stop) - delta = timedelta(minutes=1) timestamp = datetime.now() - delta @@ -77,20 +68,14 @@ def test_lambda_metric_datetime_with_extension(self): ) ] ) - self.mock_metric_extension_thread_stats.distribution.assert_not_called() @patch("datadog_lambda.metric.should_use_extension", True) def test_lambda_metric_invalid_timestamp_with_extension(self): - patcher = patch("datadog_lambda.metric.extension_thread_stats") - self.mock_metric_extension_thread_stats = patcher.start() - self.addCleanup(patcher.stop) - delta = timedelta(hours=5) timestamp = int((datetime.now() - delta).timestamp()) lambda_metric("test_timestamp", 1, timestamp) self.mock_metric_lambda_stats.distribution.assert_not_called() - self.mock_metric_extension_thread_stats.distribution.assert_not_called() def test_lambda_metric_flush_to_log(self): os.environ["DD_FLUSH_TO_LOG"] = "True" @@ -135,10 +120,6 @@ 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) @@ -217,10 +198,6 @@ def test_flush_temp_constant_tags(self): lambda_stats.thread_stats.constant_tags, original_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 7d213809bd3be14eeae23230e2857d465741c87b Mon Sep 17 00:00:00 2001 From: Aleksandr Pasechnik Date: Fri, 2 May 2025 15:05:17 -0400 Subject: [PATCH 05/10] feat: add a more explicit MetricsHandler process for routing metrics --- datadog_lambda/metric.py | 65 ++++++++++++++++++++++++++++------------ tests/test_metric.py | 38 ++++++++++++++++------- 2 files changed, 74 insertions(+), 29 deletions(-) diff --git a/datadog_lambda/metric.py b/datadog_lambda/metric.py index 79e512f8..68d3a417 100644 --- a/datadog_lambda/metric.py +++ b/datadog_lambda/metric.py @@ -3,6 +3,7 @@ # This product includes software developed at Datadog (https://www.datadoghq.com/). # Copyright 2019 Datadog, Inc. +import enum import logging import os import time @@ -15,15 +16,32 @@ logger = logging.getLogger(__name__) -lambda_stats = None -flush_in_thread = os.environ.get("DD_FLUSH_IN_THREAD", "").lower() == "true" +class MetricsHandler(enum.Enum): + EXTENSION = "extension" + FORWARDER = "forwarder" + DATADOG_API = "datadog_api" + + +def _select_metrics_handler(): + if should_use_extension: + return MetricsHandler.EXTENSION + if os.environ.get("DD_FLUSH_TO_LOG", "").lower() == "true": + return MetricsHandler.FORWARDER + return MetricsHandler.DATADOG_API + -if should_use_extension: +metrics_handler = _select_metrics_handler() +logger.debug("identified primary metrics handler as %s", metrics_handler) + + +lambda_stats = None +if metrics_handler == MetricsHandler.EXTENSION: from datadog_lambda.statsd_writer import StatsDWriter lambda_stats = StatsDWriter() -else: + +elif metrics_handler == MetricsHandler.DATADOG_API: # Periodical flushing in a background thread is NOT guaranteed to succeed # and leads to data loss. When disabled, metrics are only flushed at the # end of invocation. To make metrics submitted from a long-running Lambda @@ -31,9 +49,11 @@ from datadog_lambda.api import init_api from datadog_lambda.thread_stats_writer import ThreadStatsWriter + flush_in_thread = os.environ.get("DD_FLUSH_IN_THREAD", "").lower() == "true" init_api() lambda_stats = ThreadStatsWriter(flush_in_thread) + enhanced_metrics_enabled = ( os.environ.get("DD_ENHANCED_METRICS", "true").lower() == "true" ) @@ -44,16 +64,19 @@ def lambda_metric(metric_name, value, timestamp=None, tags=None, force_async=Fal Submit a data point to Datadog distribution metrics. https://docs.datadoghq.com/graphing/metrics/distributions/ - When DD_FLUSH_TO_LOG is True, write metric to log, and - wait for the Datadog Log Forwarder Lambda function to submit - the metrics asynchronously. + If the Datadog Lambda Extension is present, metrics are submitted to its + dogstatsd endpoint. + + When DD_FLUSH_TO_LOG is True or force_async is True, write metric to log, + and wait for the Datadog Log Forwarder Lambda function to submit the + metrics asynchronously. Otherwise, the metrics will be submitted to the Datadog API periodically and at the end of the function execution in a background thread. - Note that if the extension is present, it will override the DD_FLUSH_TO_LOG value - and always use the layer to send metrics to the extension + Note that if the extension is present, it will override the DD_FLUSH_TO_LOG + value and always use the layer to send metrics to the extension """ if not metric_name or not isinstance(metric_name, str): logger.warning( @@ -71,11 +94,10 @@ def lambda_metric(metric_name, value, timestamp=None, tags=None, force_async=Fal ) return - flush_to_logs = os.environ.get("DD_FLUSH_TO_LOG", "").lower() == "true" tags = [] if tags is None else list(tags) tags.append(dd_lambda_layer_tag) - if should_use_extension: + if metrics_handler == MetricsHandler.EXTENSION: if timestamp is not None: if isinstance(timestamp, datetime): timestamp = int(timestamp.timestamp()) @@ -94,15 +116,20 @@ def lambda_metric(metric_name, value, timestamp=None, tags=None, force_async=Fal ) lambda_stats.distribution(metric_name, value, tags=tags, timestamp=timestamp) + elif force_async or (metrics_handler == MetricsHandler.FORWARDER): + write_metric_point_to_stdout(metric_name, value, timestamp=timestamp, tags=tags) + + elif metrics_handler == MetricsHandler.DATADOG_API: + lambda_stats.distribution(metric_name, value, tags=tags, timestamp=timestamp) + else: - if flush_to_logs or force_async: - write_metric_point_to_stdout( - metric_name, value, timestamp=timestamp, tags=tags - ) - else: - lambda_stats.distribution( - metric_name, value, tags=tags, timestamp=timestamp - ) + # This should be qutie impossible, but let's at least log a message if + # it somehow happens. + logger.debug( + "Metric %s cannot be submitted because the metrics handler is not configured: %s", + metric_name, + metrics_handler, + ) def write_metric_point_to_stdout(metric_name, value, timestamp=None, tags=None): diff --git a/tests/test_metric.py b/tests/test_metric.py index d32c0876..d6d29c77 100644 --- a/tests/test_metric.py +++ b/tests/test_metric.py @@ -7,7 +7,12 @@ from datadog.api.exceptions import ClientError from datadog_lambda.api import KMS_ENCRYPTION_CONTEXT_KEY, decrypt_kms_api_key -from datadog_lambda.metric import flush_stats, lambda_metric +from datadog_lambda.metric import ( + MetricsHandler, + _select_metrics_handler, + flush_stats, + lambda_metric, +) from datadog_lambda.tags import dd_lambda_layer_tag from datadog_lambda.thread_stats_writer import ThreadStatsWriter @@ -34,15 +39,31 @@ def test_lambda_metric_tagged_with_dd_lambda_layer(self): # let's fake that the extension is present, this should override DD_FLUSH_TO_LOG @patch("datadog_lambda.metric.should_use_extension", True) - def test_lambda_metric_flush_to_log_with_extension(self): + def test_select_metrics_handler_extension_despite_flush_to_logs(self): os.environ["DD_FLUSH_TO_LOG"] = "True" + self.assertEqual(MetricsHandler.EXTENSION, _select_metrics_handler()) + del os.environ["DD_FLUSH_TO_LOG"] + + @patch("datadog_lambda.metric.should_use_extension", False) + def test_select_metrics_handler_forwarder_when_flush_to_logs(self): + os.environ["DD_FLUSH_TO_LOG"] = "True" + self.assertEqual(MetricsHandler.FORWARDER, _select_metrics_handler()) + del os.environ["DD_FLUSH_TO_LOG"] + + @patch("datadog_lambda.metric.should_use_extension", False) + def test_select_metrics_handler_dd_api_fallback(self): + os.environ["DD_FLUSH_TO_LOG"] = "False" + self.assertEqual(MetricsHandler.DATADOG_API, _select_metrics_handler()) + del os.environ["DD_FLUSH_TO_LOG"] + + @patch("datadog_lambda.metric.metrics_handler", MetricsHandler.EXTENSION) + def test_lambda_metric_flush_to_log_with_extension(self): lambda_metric("test", 1) self.mock_metric_lambda_stats.distribution.assert_has_calls( [call("test", 1, timestamp=None, tags=[dd_lambda_layer_tag])] ) - del os.environ["DD_FLUSH_TO_LOG"] - @patch("datadog_lambda.metric.should_use_extension", True) + @patch("datadog_lambda.metric.metrics_handler", MetricsHandler.EXTENSION) def test_lambda_metric_timestamp_with_extension(self): delta = timedelta(minutes=1) timestamp = int((datetime.now() - delta).timestamp()) @@ -52,7 +73,7 @@ def test_lambda_metric_timestamp_with_extension(self): [call("test_timestamp", 1, timestamp=timestamp, tags=[dd_lambda_layer_tag])] ) - @patch("datadog_lambda.metric.should_use_extension", True) + @patch("datadog_lambda.metric.metrics_handler", MetricsHandler.EXTENSION) def test_lambda_metric_datetime_with_extension(self): delta = timedelta(minutes=1) timestamp = datetime.now() - delta @@ -69,7 +90,7 @@ def test_lambda_metric_datetime_with_extension(self): ] ) - @patch("datadog_lambda.metric.should_use_extension", True) + @patch("datadog_lambda.metric.metrics_handler", MetricsHandler.EXTENSION) def test_lambda_metric_invalid_timestamp_with_extension(self): delta = timedelta(hours=5) timestamp = int((datetime.now() - delta).timestamp()) @@ -77,14 +98,11 @@ def test_lambda_metric_invalid_timestamp_with_extension(self): lambda_metric("test_timestamp", 1, timestamp) self.mock_metric_lambda_stats.distribution.assert_not_called() + @patch("datadog_lambda.metric.metrics_handler", MetricsHandler.FORWARDER) def test_lambda_metric_flush_to_log(self): - os.environ["DD_FLUSH_TO_LOG"] = "True" - lambda_metric("test", 1) self.mock_metric_lambda_stats.distribution.assert_not_called() - del os.environ["DD_FLUSH_TO_LOG"] - @patch("datadog_lambda.metric.logger.warning") def test_lambda_metric_invalid_metric_name_none(self, mock_logger_warning): lambda_metric(None, 1) From b761b28316210a6f5c65b5170708fff6a6892db9 Mon Sep 17 00:00:00 2001 From: Aleksandr Pasechnik Date: Fri, 2 May 2025 16:21:31 -0400 Subject: [PATCH 06/10] feat: new DD_LAMBDA_FIPS_MODE control for secrets and metrics --- datadog_lambda/api.py | 22 ++++++++++++++-------- datadog_lambda/fips.py | 19 +++++++++++++++++++ datadog_lambda/metric.py | 17 +++++++++++++++++ tests/test_api.py | 26 +++++++++++++++++++++++++- tests/test_metric.py | 36 ++++++++++++++++++++++++++++++++---- 5 files changed, 107 insertions(+), 13 deletions(-) create mode 100644 datadog_lambda/fips.py diff --git a/datadog_lambda/api.py b/datadog_lambda/api.py index fd3e2c17..0c978226 100644 --- a/datadog_lambda/api.py +++ b/datadog_lambda/api.py @@ -1,5 +1,7 @@ -import os import logging +import os + +from datadog_lambda.fips import enable_fips_mode logger = logging.getLogger(__name__) KMS_ENCRYPTION_CONTEXT_KEY = "LambdaFunctionName" @@ -7,9 +9,10 @@ def decrypt_kms_api_key(kms_client, ciphertext): - from botocore.exceptions import ClientError import base64 + from botocore.exceptions import ClientError + """ Decodes and deciphers the base64-encoded ciphertext given as a parameter using KMS. For this to work properly, the Lambda function must have the appropriate IAM permissions. @@ -63,10 +66,9 @@ def get_api_key() -> str: DD_API_KEY = os.environ.get("DD_API_KEY", os.environ.get("DATADOG_API_KEY", "")) LAMBDA_REGION = os.environ.get("AWS_REGION", "") - is_gov_region = LAMBDA_REGION.startswith("us-gov-") - if is_gov_region: + if enable_fips_mode: logger.debug( - "Govcloud region detected. Using FIPs endpoints for secrets management." + "FIPS mode is enabled, using FIPS endpoints for secrets management." ) if DD_API_KEY_SECRET_ARN: @@ -80,7 +82,7 @@ def get_api_key() -> str: return "" endpoint_url = ( f"https://secretsmanager-fips.{secrets_region}.amazonaws.com" - if is_gov_region + if enable_fips_mode else None ) secrets_manager_client = _boto3_client( @@ -92,7 +94,9 @@ def get_api_key() -> str: elif DD_API_KEY_SSM_NAME: # SSM endpoints: https://docs.aws.amazon.com/general/latest/gr/ssm.html fips_endpoint = ( - f"https://ssm-fips.{LAMBDA_REGION}.amazonaws.com" if is_gov_region else None + f"https://ssm-fips.{LAMBDA_REGION}.amazonaws.com" + if enable_fips_mode + else None ) ssm_client = _boto3_client("ssm", endpoint_url=fips_endpoint) api_key = ssm_client.get_parameter( @@ -101,7 +105,9 @@ def get_api_key() -> str: elif DD_KMS_API_KEY: # KMS endpoints: https://docs.aws.amazon.com/general/latest/gr/kms.html fips_endpoint = ( - f"https://kms-fips.{LAMBDA_REGION}.amazonaws.com" if is_gov_region else None + f"https://kms-fips.{LAMBDA_REGION}.amazonaws.com" + if enable_fips_mode + else None ) kms_client = _boto3_client("kms", endpoint_url=fips_endpoint) api_key = decrypt_kms_api_key(kms_client, DD_KMS_API_KEY) diff --git a/datadog_lambda/fips.py b/datadog_lambda/fips.py new file mode 100644 index 00000000..198b18ae --- /dev/null +++ b/datadog_lambda/fips.py @@ -0,0 +1,19 @@ +import logging +import os + +is_gov_region = os.environ.get("AWS_REGION", "").startswith("us-gov-") + +enable_fips_mode = ( + os.environ.get( + "DD_LAMBDA_FIPS_MODE", + "true" if is_gov_region else "false", + ).lower() + == "true" +) + +if is_gov_region or enable_fips_mode: + logger = logging.getLogger(__name__) + logger.debug( + "Python Lambda Layer FIPS mode is %s.", + "enabled" if enable_fips_mode else "not enabled", + ) diff --git a/datadog_lambda/metric.py b/datadog_lambda/metric.py index 68d3a417..d22ff976 100644 --- a/datadog_lambda/metric.py +++ b/datadog_lambda/metric.py @@ -12,6 +12,7 @@ import ujson as json from datadog_lambda.extension import should_use_extension +from datadog_lambda.fips import enable_fips_mode from datadog_lambda.tags import dd_lambda_layer_tag, get_enhanced_metrics_tags logger = logging.getLogger(__name__) @@ -21,6 +22,7 @@ class MetricsHandler(enum.Enum): EXTENSION = "extension" FORWARDER = "forwarder" DATADOG_API = "datadog_api" + NO_METRICS = "no_metrics" def _select_metrics_handler(): @@ -28,10 +30,19 @@ def _select_metrics_handler(): return MetricsHandler.EXTENSION if os.environ.get("DD_FLUSH_TO_LOG", "").lower() == "true": return MetricsHandler.FORWARDER + + if enable_fips_mode: + logger.debug( + "With FIPS mode enabled, the Datadog API metrics handler is unavailable." + ) + return MetricsHandler.NO_METRICS + return MetricsHandler.DATADOG_API metrics_handler = _select_metrics_handler() +# TODO: add a metric for this so that we can see how often the DATADOG_API +# metrics handler is actually used. logger.debug("identified primary metrics handler as %s", metrics_handler) @@ -122,6 +133,12 @@ def lambda_metric(metric_name, value, timestamp=None, tags=None, force_async=Fal elif metrics_handler == MetricsHandler.DATADOG_API: lambda_stats.distribution(metric_name, value, tags=tags, timestamp=timestamp) + elif metrics_handler == MetricsHandler.NO_METRICS: + logger.debug( + "Metric %s cannot be submitted because the metrics handler is disabled: %s", + metric_name, + ), + else: # This should be qutie impossible, but let's at least log a message if # it somehow happens. diff --git a/tests/test_api.py b/tests/test_api.py index c98d91eb..0bfd5d7f 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -1,6 +1,6 @@ import os import unittest -from unittest.mock import patch, MagicMock +from unittest.mock import MagicMock, patch import datadog_lambda.api as api @@ -22,6 +22,7 @@ def setUp(self): ) self.env_patcher.start() + @patch("datadog_lambda.api.enable_fips_mode", True) @patch("botocore.session.Session.create_client") def test_secrets_manager_fips_endpoint(self, mock_boto3_client): mock_client = MagicMock() @@ -62,6 +63,28 @@ def test_secrets_manager_different_region(self, mock_boto3_client): ) self.assertEqual(api_key, "test-api-key") + @patch("datadog_lambda.api.enable_fips_mode", True) + @patch("botocore.session.Session.create_client") + def test_secrets_manager_different_region_but_still_fips(self, mock_boto3_client): + mock_client = MagicMock() + mock_client.get_secret_value.return_value = {"SecretString": "test-api-key"} + mock_boto3_client.return_value = mock_client + + os.environ["AWS_REGION"] = "us-east-1" + os.environ[ + "DD_API_KEY_SECRET_ARN" + ] = "arn:aws:secretsmanager:us-west-1:1234567890:secret:key-name-123ABC" + + api_key = api.get_api_key() + + mock_boto3_client.assert_called_with( + "secretsmanager", + endpoint_url="https://secretsmanager-fips.us-west-1.amazonaws.com", + region_name="us-west-1", + ) + self.assertEqual(api_key, "test-api-key") + + @patch("datadog_lambda.api.enable_fips_mode", True) @patch("botocore.session.Session.create_client") def test_ssm_fips_endpoint(self, mock_boto3_client): mock_client = MagicMock() @@ -80,6 +103,7 @@ def test_ssm_fips_endpoint(self, mock_boto3_client): ) self.assertEqual(api_key, "test-api-key") + @patch("datadog_lambda.api.enable_fips_mode", True) @patch("botocore.session.Session.create_client") @patch("datadog_lambda.api.decrypt_kms_api_key") def test_kms_fips_endpoint(self, mock_decrypt_kms, mock_boto3_client): diff --git a/tests/test_metric.py b/tests/test_metric.py index d6d29c77..4edf1a79 100644 --- a/tests/test_metric.py +++ b/tests/test_metric.py @@ -19,9 +19,15 @@ class TestLambdaMetric(unittest.TestCase): def setUp(self): - patcher = patch("datadog_lambda.metric.lambda_stats") - self.mock_metric_lambda_stats = patcher.start() - self.addCleanup(patcher.stop) + lambda_stats_patcher = patch("datadog_lambda.metric.lambda_stats") + self.mock_metric_lambda_stats = lambda_stats_patcher.start() + self.addCleanup(lambda_stats_patcher.stop) + + stdout_metric_patcher = patch( + "datadog_lambda.metric.write_metric_point_to_stdout" + ) + self.mock_write_metric_point_to_stdout = stdout_metric_patcher.start() + self.addCleanup(stdout_metric_patcher.stop) def test_lambda_metric_tagged_with_dd_lambda_layer(self): lambda_metric("test", 1) @@ -56,13 +62,26 @@ def test_select_metrics_handler_dd_api_fallback(self): self.assertEqual(MetricsHandler.DATADOG_API, _select_metrics_handler()) del os.environ["DD_FLUSH_TO_LOG"] + @patch("datadog_lambda.metric.enable_fips_mode", True) + @patch("datadog_lambda.metric.should_use_extension", False) + def test_select_metrics_handler_has_no_fallback_in_fips_mode(self): + os.environ["DD_FLUSH_TO_LOG"] = "False" + self.assertEqual(MetricsHandler.NO_METRICS, _select_metrics_handler()) + del os.environ["DD_FLUSH_TO_LOG"] + @patch("datadog_lambda.metric.metrics_handler", MetricsHandler.EXTENSION) - def test_lambda_metric_flush_to_log_with_extension(self): + def test_lambda_metric_goes_to_extension_with_extension_handler(self): lambda_metric("test", 1) self.mock_metric_lambda_stats.distribution.assert_has_calls( [call("test", 1, timestamp=None, tags=[dd_lambda_layer_tag])] ) + @patch("datadog_lambda.metric.metrics_handler", MetricsHandler.NO_METRICS) + def test_lambda_metric_has_nowhere_to_go_with_no_metrics_handler(self): + lambda_metric("test", 1) + self.mock_metric_lambda_stats.distribution.assert_not_called() + self.mock_write_metric_point_to_stdout.assert_not_called() + @patch("datadog_lambda.metric.metrics_handler", MetricsHandler.EXTENSION) def test_lambda_metric_timestamp_with_extension(self): delta = timedelta(minutes=1) @@ -72,6 +91,7 @@ def test_lambda_metric_timestamp_with_extension(self): self.mock_metric_lambda_stats.distribution.assert_has_calls( [call("test_timestamp", 1, timestamp=timestamp, tags=[dd_lambda_layer_tag])] ) + self.mock_write_metric_point_to_stdout.assert_not_called() @patch("datadog_lambda.metric.metrics_handler", MetricsHandler.EXTENSION) def test_lambda_metric_datetime_with_extension(self): @@ -89,6 +109,7 @@ def test_lambda_metric_datetime_with_extension(self): ) ] ) + self.mock_write_metric_point_to_stdout.assert_not_called() @patch("datadog_lambda.metric.metrics_handler", MetricsHandler.EXTENSION) def test_lambda_metric_invalid_timestamp_with_extension(self): @@ -97,16 +118,21 @@ def test_lambda_metric_invalid_timestamp_with_extension(self): lambda_metric("test_timestamp", 1, timestamp) self.mock_metric_lambda_stats.distribution.assert_not_called() + self.mock_write_metric_point_to_stdout.assert_not_called() @patch("datadog_lambda.metric.metrics_handler", MetricsHandler.FORWARDER) def test_lambda_metric_flush_to_log(self): lambda_metric("test", 1) self.mock_metric_lambda_stats.distribution.assert_not_called() + self.mock_write_metric_point_to_stdout.assert_has_calls( + [call("test", 1, timestamp=None, tags=[dd_lambda_layer_tag])] + ) @patch("datadog_lambda.metric.logger.warning") def test_lambda_metric_invalid_metric_name_none(self, mock_logger_warning): lambda_metric(None, 1) self.mock_metric_lambda_stats.distribution.assert_not_called() + self.mock_write_metric_point_to_stdout.assert_not_called() mock_logger_warning.assert_called_once_with( "Ignoring metric submission. Invalid metric name: %s", None ) @@ -115,6 +141,7 @@ def test_lambda_metric_invalid_metric_name_none(self, mock_logger_warning): def test_lambda_metric_invalid_metric_name_not_string(self, mock_logger_warning): lambda_metric(123, 1) self.mock_metric_lambda_stats.distribution.assert_not_called() + self.mock_write_metric_point_to_stdout.assert_not_called() mock_logger_warning.assert_called_once_with( "Ignoring metric submission. Invalid metric name: %s", 123 ) @@ -123,6 +150,7 @@ def test_lambda_metric_invalid_metric_name_not_string(self, mock_logger_warning) def test_lambda_metric_non_numeric_value(self, mock_logger_warning): lambda_metric("test.non_numeric", "oops") self.mock_metric_lambda_stats.distribution.assert_not_called() + self.mock_write_metric_point_to_stdout.assert_not_called() mock_logger_warning.assert_called_once_with( "Ignoring metric submission for metric '%s' because the value is not numeric: %r", "test.non_numeric", From 360a252c8556b1b983f3df2b4d8f214191eaf805 Mon Sep 17 00:00:00 2001 From: Aleksandr Pasechnik Date: Fri, 2 May 2025 16:24:35 -0400 Subject: [PATCH 07/10] chore: remove unfortunately impossible TODO --- datadog_lambda/metric.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/datadog_lambda/metric.py b/datadog_lambda/metric.py index d22ff976..3c16c9ea 100644 --- a/datadog_lambda/metric.py +++ b/datadog_lambda/metric.py @@ -41,8 +41,6 @@ def _select_metrics_handler(): metrics_handler = _select_metrics_handler() -# TODO: add a metric for this so that we can see how often the DATADOG_API -# metrics handler is actually used. logger.debug("identified primary metrics handler as %s", metrics_handler) From f8967a790aeb28ca84f9b39a14f96e4f81d0429e Mon Sep 17 00:00:00 2001 From: Aleksandr Pasechnik Date: Fri, 2 May 2025 17:11:44 -0400 Subject: [PATCH 08/10] chore: fix typo --- datadog_lambda/metric.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datadog_lambda/metric.py b/datadog_lambda/metric.py index 3c16c9ea..1923a708 100644 --- a/datadog_lambda/metric.py +++ b/datadog_lambda/metric.py @@ -133,7 +133,7 @@ def lambda_metric(metric_name, value, timestamp=None, tags=None, force_async=Fal elif metrics_handler == MetricsHandler.NO_METRICS: logger.debug( - "Metric %s cannot be submitted because the metrics handler is disabled: %s", + "Metric %s cannot be submitted because the metrics handler is disabled", metric_name, ), From 95e2cbcc60166e73be8de2830b6505db745718ae Mon Sep 17 00:00:00 2001 From: Aleksandr Pasechnik Date: Fri, 2 May 2025 17:21:11 -0400 Subject: [PATCH 09/10] chore: add guard around lambda_stats.flush() --- datadog_lambda/metric.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/datadog_lambda/metric.py b/datadog_lambda/metric.py index 1923a708..69f0d7a7 100644 --- a/datadog_lambda/metric.py +++ b/datadog_lambda/metric.py @@ -168,7 +168,8 @@ def write_metric_point_to_stdout(metric_name, value, timestamp=None, tags=None): def flush_stats(lambda_context=None): - lambda_stats.flush() + if lambda_stats is not None: + lambda_stats.flush() def submit_enhanced_metric(metric_name, lambda_context): From f187a1b8dab0a5efb7cc045e0b4aad06f23e48e3 Mon Sep 17 00:00:00 2001 From: Aleksandr Pasechnik Date: Mon, 5 May 2025 16:00:44 -0400 Subject: [PATCH 10/10] chore: rename enable_fips_mode to fips_mode_enabled --- datadog_lambda/api.py | 10 +++++----- datadog_lambda/fips.py | 6 +++--- datadog_lambda/metric.py | 4 ++-- tests/test_api.py | 8 ++++---- tests/test_metric.py | 2 +- 5 files changed, 15 insertions(+), 15 deletions(-) diff --git a/datadog_lambda/api.py b/datadog_lambda/api.py index 0c978226..d1cee4e4 100644 --- a/datadog_lambda/api.py +++ b/datadog_lambda/api.py @@ -1,7 +1,7 @@ import logging import os -from datadog_lambda.fips import enable_fips_mode +from datadog_lambda.fips import fips_mode_enabled logger = logging.getLogger(__name__) KMS_ENCRYPTION_CONTEXT_KEY = "LambdaFunctionName" @@ -66,7 +66,7 @@ def get_api_key() -> str: DD_API_KEY = os.environ.get("DD_API_KEY", os.environ.get("DATADOG_API_KEY", "")) LAMBDA_REGION = os.environ.get("AWS_REGION", "") - if enable_fips_mode: + if fips_mode_enabled: logger.debug( "FIPS mode is enabled, using FIPS endpoints for secrets management." ) @@ -82,7 +82,7 @@ def get_api_key() -> str: return "" endpoint_url = ( f"https://secretsmanager-fips.{secrets_region}.amazonaws.com" - if enable_fips_mode + if fips_mode_enabled else None ) secrets_manager_client = _boto3_client( @@ -95,7 +95,7 @@ def get_api_key() -> str: # SSM endpoints: https://docs.aws.amazon.com/general/latest/gr/ssm.html fips_endpoint = ( f"https://ssm-fips.{LAMBDA_REGION}.amazonaws.com" - if enable_fips_mode + if fips_mode_enabled else None ) ssm_client = _boto3_client("ssm", endpoint_url=fips_endpoint) @@ -106,7 +106,7 @@ def get_api_key() -> str: # KMS endpoints: https://docs.aws.amazon.com/general/latest/gr/kms.html fips_endpoint = ( f"https://kms-fips.{LAMBDA_REGION}.amazonaws.com" - if enable_fips_mode + if fips_mode_enabled else None ) kms_client = _boto3_client("kms", endpoint_url=fips_endpoint) diff --git a/datadog_lambda/fips.py b/datadog_lambda/fips.py index 198b18ae..8442ddd9 100644 --- a/datadog_lambda/fips.py +++ b/datadog_lambda/fips.py @@ -3,7 +3,7 @@ is_gov_region = os.environ.get("AWS_REGION", "").startswith("us-gov-") -enable_fips_mode = ( +fips_mode_enabled = ( os.environ.get( "DD_LAMBDA_FIPS_MODE", "true" if is_gov_region else "false", @@ -11,9 +11,9 @@ == "true" ) -if is_gov_region or enable_fips_mode: +if is_gov_region or fips_mode_enabled: logger = logging.getLogger(__name__) logger.debug( "Python Lambda Layer FIPS mode is %s.", - "enabled" if enable_fips_mode else "not enabled", + "enabled" if fips_mode_enabled else "not enabled", ) diff --git a/datadog_lambda/metric.py b/datadog_lambda/metric.py index 69f0d7a7..0c18b517 100644 --- a/datadog_lambda/metric.py +++ b/datadog_lambda/metric.py @@ -12,7 +12,7 @@ import ujson as json from datadog_lambda.extension import should_use_extension -from datadog_lambda.fips import enable_fips_mode +from datadog_lambda.fips import fips_mode_enabled from datadog_lambda.tags import dd_lambda_layer_tag, get_enhanced_metrics_tags logger = logging.getLogger(__name__) @@ -31,7 +31,7 @@ def _select_metrics_handler(): if os.environ.get("DD_FLUSH_TO_LOG", "").lower() == "true": return MetricsHandler.FORWARDER - if enable_fips_mode: + if fips_mode_enabled: logger.debug( "With FIPS mode enabled, the Datadog API metrics handler is unavailable." ) diff --git a/tests/test_api.py b/tests/test_api.py index 0bfd5d7f..59ee4ee8 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -22,7 +22,7 @@ def setUp(self): ) self.env_patcher.start() - @patch("datadog_lambda.api.enable_fips_mode", True) + @patch("datadog_lambda.api.fips_mode_enabled", True) @patch("botocore.session.Session.create_client") def test_secrets_manager_fips_endpoint(self, mock_boto3_client): mock_client = MagicMock() @@ -63,7 +63,7 @@ def test_secrets_manager_different_region(self, mock_boto3_client): ) self.assertEqual(api_key, "test-api-key") - @patch("datadog_lambda.api.enable_fips_mode", True) + @patch("datadog_lambda.api.fips_mode_enabled", True) @patch("botocore.session.Session.create_client") def test_secrets_manager_different_region_but_still_fips(self, mock_boto3_client): mock_client = MagicMock() @@ -84,7 +84,7 @@ def test_secrets_manager_different_region_but_still_fips(self, mock_boto3_client ) self.assertEqual(api_key, "test-api-key") - @patch("datadog_lambda.api.enable_fips_mode", True) + @patch("datadog_lambda.api.fips_mode_enabled", True) @patch("botocore.session.Session.create_client") def test_ssm_fips_endpoint(self, mock_boto3_client): mock_client = MagicMock() @@ -103,7 +103,7 @@ def test_ssm_fips_endpoint(self, mock_boto3_client): ) self.assertEqual(api_key, "test-api-key") - @patch("datadog_lambda.api.enable_fips_mode", True) + @patch("datadog_lambda.api.fips_mode_enabled", True) @patch("botocore.session.Session.create_client") @patch("datadog_lambda.api.decrypt_kms_api_key") def test_kms_fips_endpoint(self, mock_decrypt_kms, mock_boto3_client): diff --git a/tests/test_metric.py b/tests/test_metric.py index 4edf1a79..a4b0be2c 100644 --- a/tests/test_metric.py +++ b/tests/test_metric.py @@ -62,7 +62,7 @@ def test_select_metrics_handler_dd_api_fallback(self): self.assertEqual(MetricsHandler.DATADOG_API, _select_metrics_handler()) del os.environ["DD_FLUSH_TO_LOG"] - @patch("datadog_lambda.metric.enable_fips_mode", True) + @patch("datadog_lambda.metric.fips_mode_enabled", True) @patch("datadog_lambda.metric.should_use_extension", False) def test_select_metrics_handler_has_no_fallback_in_fips_mode(self): os.environ["DD_FLUSH_TO_LOG"] = "False"