diff --git a/datadog_lambda/extension.py b/datadog_lambda/extension.py index 159048d7..79c0031a 100644 --- a/datadog_lambda/extension.py +++ b/datadog_lambda/extension.py @@ -1,5 +1,5 @@ import logging -from os import path +import os AGENT_URL = "http://127.0.0.1:8124" FLUSH_PATH = "/lambda/flush" @@ -9,9 +9,7 @@ def is_extension_present(): - if path.exists(EXTENSION_PATH): - return True - return False + return os.path.exists(EXTENSION_PATH) def flush_extension(): diff --git a/datadog_lambda/metric.py b/datadog_lambda/metric.py index ca23ed96..b90830e1 100644 --- a/datadog_lambda/metric.py +++ b/datadog_lambda/metric.py @@ -32,6 +32,10 @@ flush_in_thread = os.environ.get("DD_FLUSH_IN_THREAD", "").lower() == "true" lambda_stats = ThreadStatsWriter(flush_in_thread) +enhanced_metrics_enabled = ( + os.environ.get("DD_ENHANCED_METRICS", "true").lower() == "true" +) + def lambda_metric(metric_name, value, timestamp=None, tags=None, force_async=False): """ @@ -89,16 +93,6 @@ def flush_stats(): lambda_stats.flush() -def are_enhanced_metrics_enabled(): - """Check env var to find if enhanced metrics should be submitted - - Returns: - boolean for whether enhanced metrics are enabled - """ - # DD_ENHANCED_METRICS defaults to true - return os.environ.get("DD_ENHANCED_METRICS", "true").lower() == "true" - - def submit_enhanced_metric(metric_name, lambda_context): """Submits the enhanced metric with the given name @@ -106,7 +100,7 @@ def submit_enhanced_metric(metric_name, lambda_context): metric_name (str): metric name w/o enhanced prefix i.e. "invocations" or "errors" lambda_context (dict): Lambda context dict passed to the function by AWS """ - if not are_enhanced_metrics_enabled(): + if not enhanced_metrics_enabled: logger.debug( "Not submitting enhanced metric %s because enhanced metrics are disabled", metric_name, diff --git a/tests/test_wrapper.py b/tests/test_wrapper.py index bb26b661..525416ff 100644 --- a/tests/test_wrapper.py +++ b/tests/test_wrapper.py @@ -465,7 +465,9 @@ def lambda_handler(event, context): ) def test_no_enhanced_metrics_without_env_var(self): - os.environ["DD_ENHANCED_METRICS"] = "false" + patcher = patch("datadog_lambda.metric.enhanced_metrics_enabled", False) + patcher.start() + self.addCleanup(patcher.stop) @wrapper.datadog_lambda_wrapper def lambda_handler(event, context): @@ -478,8 +480,6 @@ def lambda_handler(event, context): self.mock_write_metric_point_to_stdout.assert_not_called() - del os.environ["DD_ENHANCED_METRICS"] - def test_only_one_wrapper_in_use(self): patcher = patch("datadog_lambda.wrapper.submit_invocations_metric") self.mock_submit_invocations_metric = patcher.start()