diff --git a/datadog_lambda/api.py b/datadog_lambda/api.py index d1cee4e4..4921dae9 100644 --- a/datadog_lambda/api.py +++ b/datadog_lambda/api.py @@ -1,7 +1,7 @@ import logging import os -from datadog_lambda.fips import fips_mode_enabled +from datadog_lambda.config import config logger = logging.getLogger(__name__) KMS_ENCRYPTION_CONTEXT_KEY = "LambdaFunctionName" @@ -29,7 +29,6 @@ def decrypt_kms_api_key(kms_client, ciphertext): is added. We need to try decrypting the API key both with and without the encryption context. """ # Try without encryption context, in case API key was encrypted using the AWS CLI - function_name = os.environ.get("AWS_LAMBDA_FUNCTION_NAME") try: plaintext = kms_client.decrypt(CiphertextBlob=decoded_bytes)[ "Plaintext" @@ -43,7 +42,7 @@ def decrypt_kms_api_key(kms_client, ciphertext): plaintext = kms_client.decrypt( CiphertextBlob=decoded_bytes, EncryptionContext={ - KMS_ENCRYPTION_CONTEXT_KEY: function_name, + KMS_ENCRYPTION_CONTEXT_KEY: config.function_name, }, )["Plaintext"].decode("utf-8") @@ -66,7 +65,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 fips_mode_enabled: + if config.fips_mode_enabled: logger.debug( "FIPS mode is enabled, using FIPS endpoints for secrets management." ) @@ -82,7 +81,7 @@ def get_api_key() -> str: return "" endpoint_url = ( f"https://secretsmanager-fips.{secrets_region}.amazonaws.com" - if fips_mode_enabled + if config.fips_mode_enabled else None ) secrets_manager_client = _boto3_client( @@ -95,7 +94,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 fips_mode_enabled + if config.fips_mode_enabled else None ) ssm_client = _boto3_client("ssm", endpoint_url=fips_endpoint) @@ -106,7 +105,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 fips_mode_enabled + if config.fips_mode_enabled else None ) kms_client = _boto3_client("kms", endpoint_url=fips_endpoint) @@ -118,7 +117,7 @@ def get_api_key() -> str: def init_api(): - if not os.environ.get("DD_FLUSH_TO_LOG", "").lower() == "true": + if not config.flush_to_log: # 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 with the extension unless sending metrics with timestamps diff --git a/datadog_lambda/cold_start.py b/datadog_lambda/cold_start.py index ea10ea20..96cb2074 100644 --- a/datadog_lambda/cold_start.py +++ b/datadog_lambda/cold_start.py @@ -1,8 +1,9 @@ import time -import os from typing import List, Hashable import logging +from datadog_lambda.config import config + logger = logging.getLogger(__name__) _cold_start = True @@ -86,14 +87,12 @@ def reset_node_stacks(): def push_node(module_name, file_path): node = ImportNode(module_name, file_path, time.time_ns()) - global import_stack if import_stack: import_stack[-1].children.append(node) import_stack.append(node) def pop_node(module_name): - global import_stack if not import_stack: return node = import_stack.pop() @@ -102,7 +101,6 @@ def pop_node(module_name): end_time_ns = time.time_ns() node.end_time_ns = end_time_ns if not import_stack: # import_stack empty, a root node has been found - global root_nodes root_nodes.append(node) @@ -147,11 +145,7 @@ def wrapped_find_spec(*args, **kwargs): def initialize_cold_start_tracing(): - if ( - is_new_sandbox() - and os.environ.get("DD_TRACE_ENABLED", "true").lower() == "true" - and os.environ.get("DD_COLD_START_TRACING", "true").lower() == "true" - ): + if is_new_sandbox() and config.trace_enabled and config.cold_start_tracing: from sys import meta_path for importer in meta_path: diff --git a/datadog_lambda/config.py b/datadog_lambda/config.py new file mode 100644 index 00000000..881993a4 --- /dev/null +++ b/datadog_lambda/config.py @@ -0,0 +1,70 @@ +# Unless explicitly stated otherwise all files in this repository are licensed +# under the Apache License Version 2.0. +# This product includes software developed at Datadog (https://www.datadoghq.com/). +# Copyright 2019 Datadog, Inc. + +import logging +import os + + +def _get_env(key, default=None, cast=None): + @property + def _getter(self): + if not hasattr(self, prop_key): + val = os.environ.get(key, default) + if cast is not None: + val = cast(val) + setattr(self, prop_key, val) + return getattr(self, prop_key) + + prop_key = f"_config_{key}" + return _getter + + +def as_bool(val): + return val.lower() == "true" or val == "1" + + +class Config: + + add_span_pointers = _get_env("DD_BOTOCORE_ADD_SPAN_POINTERS", "true", as_bool) + cold_start_tracing = _get_env("DD_COLD_START_TRACING", "true", as_bool) + enhanced_metrics_enabled = _get_env("DD_ENHANCED_METRICS", "true", as_bool) + flush_in_thread = _get_env("DD_FLUSH_IN_THREAD", "false", as_bool) + flush_to_log = _get_env("DD_FLUSH_TO_LOG", "false", as_bool) + function_name = _get_env("AWS_LAMBDA_FUNCTION_NAME") + is_gov_region = _get_env("AWS_REGION", "", lambda x: x.startswith("us-gov-")) + is_in_tests = _get_env("DD_INTEGRATION_TEST", "false", as_bool) + is_lambda_context = _get_env("AWS_LAMBDA_FUNCTION_NAME", None, bool) + otel_enabled = _get_env("DD_TRACE_OTEL_ENABLED", "false", as_bool) + telemetry_enabled = _get_env( + "DD_INSTRUMENTATION_TELEMETRY_ENABLED", "false", as_bool + ) + trace_enabled = _get_env("DD_TRACE_ENABLED", "true", as_bool) + + @property + def fips_mode_enabled(self): + if not hasattr(self, "_config_fips_mode_enabled"): + self._config_fips_mode_enabled = ( + os.environ.get( + "DD_LAMBDA_FIPS_MODE", + "true" if self.is_gov_region else "false", + ).lower() + == "true" + ) + return self._config_fips_mode_enabled + + def reset(self): + for attr in dir(self): + if attr.startswith("_config_"): + delattr(self, attr) + + +config = Config() + +if config.is_gov_region or config.fips_mode_enabled: + logger = logging.getLogger(__name__) + logger.debug( + "Python Lambda Layer FIPS mode is %s.", + "enabled" if config.fips_mode_enabled else "not enabled", + ) diff --git a/datadog_lambda/fips.py b/datadog_lambda/fips.py deleted file mode 100644 index 8442ddd9..00000000 --- a/datadog_lambda/fips.py +++ /dev/null @@ -1,19 +0,0 @@ -import logging -import os - -is_gov_region = os.environ.get("AWS_REGION", "").startswith("us-gov-") - -fips_mode_enabled = ( - os.environ.get( - "DD_LAMBDA_FIPS_MODE", - "true" if is_gov_region else "false", - ).lower() - == "true" -) - -if is_gov_region or fips_mode_enabled: - logger = logging.getLogger(__name__) - logger.debug( - "Python Lambda Layer FIPS mode is %s.", - "enabled" if fips_mode_enabled else "not enabled", - ) diff --git a/datadog_lambda/metric.py b/datadog_lambda/metric.py index c9b978d6..73bbeca3 100644 --- a/datadog_lambda/metric.py +++ b/datadog_lambda/metric.py @@ -5,14 +5,13 @@ import enum import logging -import os import time from datetime import datetime, timedelta import ujson as json +from datadog_lambda.config import config from datadog_lambda.extension import should_use_extension -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__) @@ -28,10 +27,10 @@ class MetricsHandler(enum.Enum): def _select_metrics_handler(): if should_use_extension: return MetricsHandler.EXTENSION - if os.environ.get("DD_FLUSH_TO_LOG", "").lower() == "true": + if config.flush_to_log: return MetricsHandler.FORWARDER - if fips_mode_enabled: + if config.fips_mode_enabled: logger.debug( "With FIPS mode enabled, the Datadog API metrics handler is unavailable." ) @@ -58,14 +57,8 @@ def _select_metrics_handler(): 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" -) + lambda_stats = ThreadStatsWriter(config.flush_in_thread) def lambda_metric(metric_name, value, timestamp=None, tags=None, force_async=False): @@ -191,7 +184,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 (object): Lambda context dict passed to the function by AWS """ - if not enhanced_metrics_enabled: + if not config.enhanced_metrics_enabled: logger.debug( "Not submitting enhanced metric %s because enhanced metrics are disabled", metric_name, diff --git a/datadog_lambda/patch.py b/datadog_lambda/patch.py index 5b8a92c5..e5aa5540 100644 --- a/datadog_lambda/patch.py +++ b/datadog_lambda/patch.py @@ -3,7 +3,6 @@ # This product includes software developed at Datadog (https://www.datadoghq.com/). # Copyright 2019 Datadog, Inc. -import os import sys import logging import zlib @@ -13,10 +12,8 @@ from wrapt.importer import when_imported from ddtrace import patch_all as patch_all_dd -from datadog_lambda.tracing import ( - get_dd_trace_context, - dd_tracing_enabled, -) +from datadog_lambda import config +from datadog_lambda.tracing import get_dd_trace_context from collections.abc import MutableMapping logger = logging.getLogger(__name__) @@ -32,7 +29,7 @@ def patch_all(): """ _patch_for_integration_tests() - if dd_tracing_enabled: + if config.trace_enabled: patch_all_dd() else: _patch_http() @@ -44,8 +41,7 @@ def _patch_for_integration_tests(): Patch `requests` to log the outgoing requests for integration tests. """ global _integration_tests_patched - is_in_tests = os.environ.get("DD_INTEGRATION_TEST", "false").lower() == "true" - if not _integration_tests_patched and is_in_tests: + if not _integration_tests_patched and config.is_in_tests: wrap("requests", "Session.send", _log_request) _integration_tests_patched = True diff --git a/datadog_lambda/span_pointers.py b/datadog_lambda/span_pointers.py index 40d959e6..45925d92 100644 --- a/datadog_lambda/span_pointers.py +++ b/datadog_lambda/span_pointers.py @@ -1,12 +1,12 @@ from itertools import chain import logging -import os from typing import List from typing import Optional from ddtrace._trace._span_pointer import _SpanPointerDirection from ddtrace._trace._span_pointer import _SpanPointerDescription +from datadog_lambda.config import config from datadog_lambda.metric import submit_dynamodb_stream_type_metric from datadog_lambda.trigger import EventTypes @@ -14,15 +14,10 @@ logger = logging.getLogger(__name__) -dd_botocore_add_span_pointers = os.environ.get( - "DD_BOTOCORE_ADD_SPAN_POINTERS", "true" -).lower() in ("true", "1") - - def calculate_span_pointers( event_source, event, - botocore_add_span_pointers=dd_botocore_add_span_pointers, + botocore_add_span_pointers=config.add_span_pointers, ) -> List[_SpanPointerDescription]: try: if botocore_add_span_pointers: diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 9a27673c..c819a6e4 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -32,6 +32,8 @@ from ddtrace import __version__ as ddtrace_version from ddtrace.propagation.http import HTTPPropagator from ddtrace.trace import Context, Span, tracer + +from datadog_lambda.config import config from datadog_lambda import __version__ as datadog_lambda_version from datadog_lambda.trigger import ( _EventSource, @@ -42,10 +44,7 @@ EventSubtypes, ) -dd_trace_otel_enabled = ( - os.environ.get("DD_TRACE_OTEL_ENABLED", "false").lower() == "true" -) -if dd_trace_otel_enabled: +if config.otel_enabled: from opentelemetry.trace import set_tracer_provider from ddtrace.opentelemetry import TracerProvider @@ -55,18 +54,11 @@ logger = logging.getLogger(__name__) dd_trace_context = None -dd_tracing_enabled = os.environ.get("DD_TRACE_ENABLED", "false").lower() == "true" -if dd_tracing_enabled: +if config.trace_enabled and config.telemetry_enabled: # Enable the telemetry client if the user has opted in - if ( - os.environ.get("DD_INSTRUMENTATION_TELEMETRY_ENABLED", "false").lower() - == "true" - ): - from ddtrace.internal.telemetry import telemetry_writer + from ddtrace.internal.telemetry import telemetry_writer - telemetry_writer.enable() - -is_lambda_context = os.environ.get(XrayDaemon.FUNCTION_NAME_HEADER_NAME) != "" + telemetry_writer.enable() propagator = HTTPPropagator() @@ -97,7 +89,7 @@ def _convert_xray_sampling(xray_sampled): def _get_xray_trace_context(): - if not is_lambda_context: + if not config.is_lambda_context: return None xray_trace_entity = parse_xray_header( @@ -639,13 +631,11 @@ def get_dd_trace_context_obj(): automatically, but this function can be used to manually inject the trace context to an outgoing request. """ - if dd_tracing_enabled: + if config.trace_enabled: dd_trace_py_context = _get_dd_trace_py_context() if _is_context_complete(dd_trace_py_context): return dd_trace_py_context - global dd_trace_context - try: xray_context = _get_xray_trace_context() # xray (sub)segment except Exception as e: @@ -690,10 +680,10 @@ def set_correlation_ids(): TODO: Remove me when Datadog tracer is natively supported in Lambda. """ - if not is_lambda_context: + if not config.is_lambda_context: logger.debug("set_correlation_ids is only supported in LambdaContext") return - if dd_tracing_enabled: + if config.trace_enabled: logger.debug("using ddtrace implementation for spans") return @@ -1476,7 +1466,7 @@ def emit_telemetry_on_exception_outside_of_handler( Emit an enhanced error metric and create a span for exceptions occurring outside the handler """ submit_errors_metric(None) - if dd_tracing_enabled: + if config.trace_enabled: span = tracer.trace( "aws.lambda", service="aws.lambda", diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index e5460118..92be5ef5 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -17,6 +17,7 @@ is_new_sandbox, ColdStartTracer, ) +from datadog_lambda.config import config from datadog_lambda.constants import ( TraceContextSource, XraySubsegment, @@ -29,7 +30,6 @@ extract_dd_trace_context, create_dd_dummy_metadata_subsegment, inject_correlation_ids, - dd_tracing_enabled, mark_trace_as_error_for_5xx_responses, set_correlation_ids, set_dd_trace_py_root, @@ -172,7 +172,7 @@ def __init__(self, func): self.span = None self.inferred_span = None depends_on_dd_tracing_enabled = ( - lambda original_boolean: dd_tracing_enabled and original_boolean + lambda original_boolean: config.trace_enabled and original_boolean ) self.make_inferred_span = depends_on_dd_tracing_enabled( os.environ.get(DD_TRACE_MANAGED_SERVICES, "true").lower() == "true" @@ -315,7 +315,7 @@ def _before(self, event, context): XraySubsegment.TRACE_KEY, ) - if dd_tracing_enabled: + if config.trace_enabled: set_dd_trace_py_root(trace_context_source, self.merge_xray_traces) if self.make_inferred_span: self.inferred_span = create_inferred_span( diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 00000000..091265e9 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,8 @@ +import pytest + +from datadog_lambda.config import config + + +@pytest.fixture(autouse=True) +def reset_config(): + config.reset() diff --git a/tests/test_api.py b/tests/test_api.py index 59ee4ee8..7fcc3c22 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -22,7 +22,10 @@ def setUp(self): ) self.env_patcher.start() - @patch("datadog_lambda.api.fips_mode_enabled", True) + def tearDown(self): + del os.environ["AWS_REGION"] + + @patch("datadog_lambda.config.Config.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 +66,7 @@ def test_secrets_manager_different_region(self, mock_boto3_client): ) self.assertEqual(api_key, "test-api-key") - @patch("datadog_lambda.api.fips_mode_enabled", True) + @patch("datadog_lambda.config.Config.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 +87,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.fips_mode_enabled", True) + @patch("datadog_lambda.config.Config.fips_mode_enabled", True) @patch("botocore.session.Session.create_client") def test_ssm_fips_endpoint(self, mock_boto3_client): mock_client = MagicMock() @@ -103,7 +106,7 @@ def test_ssm_fips_endpoint(self, mock_boto3_client): ) self.assertEqual(api_key, "test-api-key") - @patch("datadog_lambda.api.fips_mode_enabled", True) + @patch("datadog_lambda.config.Config.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_cold_start.py b/tests/test_cold_start.py index c7444c49..f90df270 100644 --- a/tests/test_cold_start.py +++ b/tests/test_cold_start.py @@ -247,7 +247,7 @@ def finish(span): monkeypatch.setattr(wrapper.tracer, "_on_span_finish", finish) monkeypatch.setattr(wrapper, "is_new_sandbox", lambda: True) - monkeypatch.setattr("datadog_lambda.wrapper.dd_tracing_enabled", True) + monkeypatch.setattr("datadog_lambda.config.Config.trace_enabled", True) monkeypatch.setenv( "DD_COLD_START_TRACE_SKIP_LIB", "ddtrace.contrib.logging,datadog_lambda.wrapper" ) diff --git a/tests/test_config.py b/tests/test_config.py new file mode 100644 index 00000000..868f5a33 --- /dev/null +++ b/tests/test_config.py @@ -0,0 +1,158 @@ +import pytest + +from datadog_lambda.config import config + + +@pytest.fixture +def setenv(monkeypatch): + def set_env(key, value): + if value is None: + monkeypatch.delenv(key, raising=False) + else: + monkeypatch.setenv(key, value) + + return set_env + + +_test_config_from_environ = ( + ("AWS_LAMBDA_FUNCTION_NAME", "function_name", None, None), + ("AWS_LAMBDA_FUNCTION_NAME", "function_name", "", ""), + ("AWS_LAMBDA_FUNCTION_NAME", "function_name", "my_function", "my_function"), + ("DD_FLUSH_TO_LOG", "flush_to_log", None, False), + ("DD_FLUSH_TO_LOG", "flush_to_log", "", False), + ("DD_FLUSH_TO_LOG", "flush_to_log", "true", True), + ("DD_FLUSH_TO_LOG", "flush_to_log", "TRUE", True), + ("DD_FLUSH_TO_LOG", "flush_to_log", "false", False), + ("DD_FLUSH_TO_LOG", "flush_to_log", "FALSE", False), + ("DD_FLUSH_TO_LOG", "flush_to_log", "1", True), # CHANGED + ("DD_FLUSH_TO_LOG", "flush_to_log", "0", False), + ("DD_FLUSH_TO_LOG", "flush_to_log", "purple", False), + ("DD_TRACE_ENABLED", "trace_enabled", None, True), + ("DD_TRACE_ENABLED", "trace_enabled", "", False), + ("DD_TRACE_ENABLED", "trace_enabled", "true", True), + ("DD_TRACE_ENABLED", "trace_enabled", "TRUE", True), + ("DD_TRACE_ENABLED", "trace_enabled", "false", False), + ("DD_TRACE_ENABLED", "trace_enabled", "FALSE", False), + ("DD_TRACE_ENABLED", "trace_enabled", "1", True), # CHANGED + ("DD_TRACE_ENABLED", "trace_enabled", "0", False), + ("DD_TRACE_ENABLED", "trace_enabled", "purple", False), + ("DD_COLD_START_TRACING", "cold_start_tracing", None, True), + ("DD_COLD_START_TRACING", "cold_start_tracing", "", False), + ("DD_COLD_START_TRACING", "cold_start_tracing", "true", True), + ("DD_COLD_START_TRACING", "cold_start_tracing", "TRUE", True), + ("DD_COLD_START_TRACING", "cold_start_tracing", "false", False), + ("DD_COLD_START_TRACING", "cold_start_tracing", "FALSE", False), + ("DD_COLD_START_TRACING", "cold_start_tracing", "1", True), # CHANGED + ("DD_COLD_START_TRACING", "cold_start_tracing", "0", False), + ("DD_COLD_START_TRACING", "cold_start_tracing", "purple", False), + ("AWS_REGION", "is_gov_region", None, False), + ("AWS_REGION", "is_gov_region", "", False), + ("AWS_REGION", "is_gov_region", "us-gov-1", True), + ("AWS_REGION", "is_gov_region", "us-est-1", False), + ("DD_FLUSH_IN_THREAD", "flush_in_thread", None, False), + ("DD_FLUSH_IN_THREAD", "flush_in_thread", "", False), + ("DD_FLUSH_IN_THREAD", "flush_in_thread", "true", True), + ("DD_FLUSH_IN_THREAD", "flush_in_thread", "TRUE", True), + ("DD_FLUSH_IN_THREAD", "flush_in_thread", "false", False), + ("DD_FLUSH_IN_THREAD", "flush_in_thread", "FALSE", False), + ("DD_FLUSH_IN_THREAD", "flush_in_thread", "1", True), # CHANGED + ("DD_FLUSH_IN_THREAD", "flush_in_thread", "0", False), + ("DD_FLUSH_IN_THREAD", "flush_in_thread", "purple", False), + ("DD_ENHANCED_METRICS", "enhanced_metrics_enabled", None, True), + ("DD_ENHANCED_METRICS", "enhanced_metrics_enabled", "", False), + ("DD_ENHANCED_METRICS", "enhanced_metrics_enabled", "true", True), + ("DD_ENHANCED_METRICS", "enhanced_metrics_enabled", "TRUE", True), + ("DD_ENHANCED_METRICS", "enhanced_metrics_enabled", "false", False), + ("DD_ENHANCED_METRICS", "enhanced_metrics_enabled", "FALSE", False), + ("DD_ENHANCED_METRICS", "enhanced_metrics_enabled", "1", True), # CHANGED + ("DD_ENHANCED_METRICS", "enhanced_metrics_enabled", "0", False), + ("DD_ENHANCED_METRICS", "enhanced_metrics_enabled", "purple", False), + ("DD_INTEGRATION_TEST", "is_in_tests", None, False), + ("DD_INTEGRATION_TEST", "is_in_tests", "", False), + ("DD_INTEGRATION_TEST", "is_in_tests", "true", True), + ("DD_INTEGRATION_TEST", "is_in_tests", "TRUE", True), + ("DD_INTEGRATION_TEST", "is_in_tests", "false", False), + ("DD_INTEGRATION_TEST", "is_in_tests", "FALSE", False), + ("DD_INTEGRATION_TEST", "is_in_tests", "1", True), # CHANGED + ("DD_INTEGRATION_TEST", "is_in_tests", "0", False), + ("DD_INTEGRATION_TEST", "is_in_tests", "purple", False), + ("DD_BOTOCORE_ADD_SPAN_POINTERS", "add_span_pointers", None, True), + ("DD_BOTOCORE_ADD_SPAN_POINTERS", "add_span_pointers", "", False), + ("DD_BOTOCORE_ADD_SPAN_POINTERS", "add_span_pointers", "true", True), + ("DD_BOTOCORE_ADD_SPAN_POINTERS", "add_span_pointers", "TRUE", True), + ("DD_BOTOCORE_ADD_SPAN_POINTERS", "add_span_pointers", "false", False), + ("DD_BOTOCORE_ADD_SPAN_POINTERS", "add_span_pointers", "FALSE", False), + ("DD_BOTOCORE_ADD_SPAN_POINTERS", "add_span_pointers", "1", True), + ("DD_BOTOCORE_ADD_SPAN_POINTERS", "add_span_pointers", "0", False), + ("DD_BOTOCORE_ADD_SPAN_POINTERS", "add_span_pointers", "purple", False), + ("DD_TRACE_OTEL_ENABLED", "otel_enabled", None, False), + ("DD_TRACE_OTEL_ENABLED", "otel_enabled", "", False), + ("DD_TRACE_OTEL_ENABLED", "otel_enabled", "true", True), + ("DD_TRACE_OTEL_ENABLED", "otel_enabled", "TRUE", True), + ("DD_TRACE_OTEL_ENABLED", "otel_enabled", "false", False), + ("DD_TRACE_OTEL_ENABLED", "otel_enabled", "FALSE", False), + ("DD_TRACE_OTEL_ENABLED", "otel_enabled", "1", True), # CHANGED + ("DD_TRACE_OTEL_ENABLED", "otel_enabled", "0", False), + ("DD_TRACE_OTEL_ENABLED", "otel_enabled", "purple", False), + ("AWS_LAMBDA_FUNCTION_NAME", "is_lambda_context", None, False), + ("AWS_LAMBDA_FUNCTION_NAME", "is_lambda_context", "", False), + ("AWS_LAMBDA_FUNCTION_NAME", "is_lambda_context", "my_function", True), + ("DD_INSTRUMENTATION_TELEMETRY_ENABLED", "telemetry_enabled", None, False), + ("DD_INSTRUMENTATION_TELEMETRY_ENABLED", "telemetry_enabled", "", False), + ("DD_INSTRUMENTATION_TELEMETRY_ENABLED", "telemetry_enabled", "true", True), + ("DD_INSTRUMENTATION_TELEMETRY_ENABLED", "telemetry_enabled", "TRUE", True), + ("DD_INSTRUMENTATION_TELEMETRY_ENABLED", "telemetry_enabled", "false", False), + ("DD_INSTRUMENTATION_TELEMETRY_ENABLED", "telemetry_enabled", "FALSE", False), + ("DD_INSTRUMENTATION_TELEMETRY_ENABLED", "telemetry_enabled", "1", True), # CHANGED + ("DD_INSTRUMENTATION_TELEMETRY_ENABLED", "telemetry_enabled", "0", False), + ("DD_INSTRUMENTATION_TELEMETRY_ENABLED", "telemetry_enabled", "purple", False), +) + + +@pytest.mark.parametrize("env_key,conf_key,env_val,conf_val", _test_config_from_environ) +def test_config_from_environ(env_key, conf_key, env_val, conf_val, setenv): + setenv(env_key, env_val) + assert getattr(config, conf_key) == conf_val + + +_test_fips_mode_from_environ = ( + (None, None, False), + (None, "", False), + (None, "us-gov-1", True), + (None, "us-east-1", False), + ("", None, False), + ("", "", False), + ("", "us-gov-1", False), + ("", "us-east-1", False), + ("true", None, True), + ("true", "", True), + ("true", "us-gov-1", True), + ("true", "us-east-1", True), + ("TRUE", None, True), + ("TRUE", "", True), + ("TRUE", "us-gov-1", True), + ("TRUE", "us-east-1", True), + ("false", None, False), + ("false", "", False), + ("false", "us-gov-1", False), + ("false", "us-east-1", False), + ("FALSE", None, False), + ("FALSE", "", False), + ("FALSE", "us-gov-1", False), + ("FALSE", "us-east-1", False), + ("1", None, False), + ("1", "", False), + ("1", "us-gov-1", False), + ("1", "us-east-1", False), + ("0", None, False), + ("0", "", False), + ("0", "us-gov-1", False), + ("0", "us-east-1", False), +) + + +@pytest.mark.parametrize("fips_mode,region,conf_val", _test_fips_mode_from_environ) +def test_fips_mode_from_environ(fips_mode, region, conf_val, setenv): + setenv("DD_LAMBDA_FIPS_MODE", fips_mode) + setenv("AWS_REGION", region) + assert config.fips_mode_enabled == conf_val diff --git a/tests/test_metric.py b/tests/test_metric.py index e7dab2c3..3de43334 100644 --- a/tests/test_metric.py +++ b/tests/test_metric.py @@ -7,6 +7,7 @@ from datadog.api.exceptions import ClientError from datadog_lambda.api import KMS_ENCRYPTION_CONTEXT_KEY, decrypt_kms_api_key +from datadog_lambda.config import config from datadog_lambda.metric import ( MetricsHandler, _select_metrics_handler, @@ -19,6 +20,7 @@ class TestLambdaMetric(unittest.TestCase): def setUp(self): + config.reset() lambda_stats_patcher = patch("datadog_lambda.metric.lambda_stats") self.mock_metric_lambda_stats = lambda_stats_patcher.start() self.addCleanup(lambda_stats_patcher.stop) @@ -62,7 +64,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.fips_mode_enabled", True) + @patch("datadog_lambda.config.Config.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" diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 0a961a62..197c557d 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -15,6 +15,7 @@ from ddtrace._trace._span_pointer import _SpanPointerDirection from ddtrace._trace._span_pointer import _SpanPointerDescription +from datadog_lambda.config import config from datadog_lambda.constants import ( SamplingPriority, TraceHeader, @@ -251,20 +252,16 @@ def test_extract_dd_trace_context(event, expect): class TestExtractAndGetDDTraceContext(unittest.TestCase): def setUp(self): - global dd_tracing_enabled - dd_tracing_enabled = False os.environ["_X_AMZN_TRACE_ID"] = fake_xray_header_value patcher = patch("datadog_lambda.tracing.send_segment") self.mock_send_segment = patcher.start() self.addCleanup(patcher.stop) - patcher = patch("datadog_lambda.tracing.is_lambda_context") + patcher = patch("datadog_lambda.config.Config.is_lambda_context") self.mock_is_lambda_context = patcher.start() self.mock_is_lambda_context.return_value = True self.addCleanup(patcher.stop) def tearDown(self): - global dd_tracing_enabled - dd_tracing_enabled = False del os.environ["_X_AMZN_TRACE_ID"] @with_trace_propagation_style("datadog") @@ -975,6 +972,7 @@ def test_convert_xray_sampling(self): class TestLogsInjection(unittest.TestCase): def setUp(self): + config.reset() patcher = patch("datadog_lambda.tracing.get_dd_trace_context_obj") self.mock_get_dd_trace_context = patcher.start() self.mock_get_dd_trace_context.return_value = Context( @@ -984,11 +982,12 @@ def setUp(self): ) self.addCleanup(patcher.stop) - patcher = patch("datadog_lambda.tracing.is_lambda_context") + patcher = patch("datadog_lambda.config.Config.is_lambda_context") self.mock_is_lambda_context = patcher.start() self.mock_is_lambda_context.return_value = True self.addCleanup(patcher.stop) + @patch("datadog_lambda.config.Config.trace_enabled", False) def test_set_correlation_ids(self): set_correlation_ids() span = tracer.current_span() @@ -1124,13 +1123,11 @@ def test_function_with_span_pointers(self): class TestSetTraceRootSpan(unittest.TestCase): def setUp(self): - global dd_tracing_enabled - dd_tracing_enabled = False os.environ["_X_AMZN_TRACE_ID"] = fake_xray_header_value patcher = patch("datadog_lambda.tracing.send_segment") self.mock_send_segment = patcher.start() self.addCleanup(patcher.stop) - patcher = patch("datadog_lambda.tracing.is_lambda_context") + patcher = patch("datadog_lambda.config.Config.is_lambda_context") self.mock_is_lambda_context = patcher.start() self.mock_is_lambda_context.return_value = True self.addCleanup(patcher.stop) @@ -1143,8 +1140,6 @@ def setUp(self): self.addCleanup(patcher.stop) def tearDown(self): - global dd_tracing_enabled - dd_tracing_enabled = False del os.environ["_X_AMZN_TRACE_ID"] def test_mixed_parent_context_when_merging(self): @@ -1245,6 +1240,7 @@ def test_get_service_mapping(self): create_service_mapping(os.environ["DD_SERVICE_MAPPING"]) ) self.assertEqual(self.get_service_mapping(), expected_output) + del os.environ["DD_SERVICE_MAPPING"] def test_set_service_mapping(self): new_service_mapping = {"api3": "service3", "api4": "service4"} @@ -1285,6 +1281,8 @@ def test_determine_service_name(self): "default", ) + del os.environ["DD_SERVICE_MAPPING"] + def test_remaps_all_inferred_span_service_names_from_api_gateway_event(self): new_service_mapping = {"lambda_api_gateway": "new-name"} self.set_service_mapping(new_service_mapping) @@ -2386,7 +2384,7 @@ def test_deterministic_m5_hash__always_leading_with_zero(self): class TestExceptionOutsideHandler(unittest.TestCase): - @patch("datadog_lambda.tracing.dd_tracing_enabled", True) + @patch("datadog_lambda.config.Config.trace_enabled", True) @patch("datadog_lambda.tracing.submit_errors_metric") @patch("time.time_ns", return_value=42) def test_exception_outside_handler_tracing_enabled( @@ -2427,7 +2425,7 @@ def test_exception_outside_handler_tracing_enabled( assert mock_span.error == 1 assert mock_span.start_ns == 42 - @patch("datadog_lambda.tracing.dd_tracing_enabled", False) + @patch("datadog_lambda.config.Config.trace_enabled", False) @patch("datadog_lambda.tracing.submit_errors_metric") @patch("time.time_ns", return_value=42) def test_exception_outside_handler_tracing_disabled( diff --git a/tests/test_wrapper.py b/tests/test_wrapper.py index 4b243036..8bb3011d 100644 --- a/tests/test_wrapper.py +++ b/tests/test_wrapper.py @@ -8,6 +8,8 @@ import datadog_lambda.wrapper as wrapper import datadog_lambda.xray as xray + +from datadog_lambda.config import config from datadog_lambda.metric import lambda_metric from datadog_lambda.thread_stats_writer import ThreadStatsWriter from ddtrace.trace import Span, tracer @@ -24,7 +26,6 @@ def setUp(self): patch("ddtrace.internal.writer.AgentWriter.flush_queue").start() wrapper.datadog_lambda_wrapper._force_wrap = True - wrapper.dd_tracing_enabled = True patcher = patch( "datadog.threadstats.reporters.HttpReporter.flush_distributions" ) @@ -76,9 +77,8 @@ def setUp(self): self.mock_dd_lambda_layer_tag = patcher.start() self.addCleanup(patcher.stop) + @patch("datadog_lambda.config.Config.trace_enabled", False) def test_datadog_lambda_wrapper(self): - wrapper.dd_tracing_enabled = False - @wrapper.datadog_lambda_wrapper def lambda_handler(event, context): lambda_metric("test.metric", 100) @@ -88,7 +88,6 @@ def lambda_handler(event, context): lambda_context = get_mock_context() lambda_handler(lambda_event, lambda_context) - wrapper.dd_tracing_enabled = True self.mock_threadstats_flush_distributions.assert_has_calls( [ call( @@ -185,9 +184,9 @@ def lambda_handler(event, context): metric_module.lambda_stats.stop() metric_module.lambda_stats = ThreadStatsWriter(False) + @patch("datadog_lambda.config.Config.trace_enabled", False) def test_datadog_lambda_wrapper_inject_correlation_ids(self): os.environ["DD_LOGS_INJECTION"] = "True" - wrapper.dd_tracing_enabled = False @wrapper.datadog_lambda_wrapper def lambda_handler(event, context): @@ -195,7 +194,6 @@ def lambda_handler(event, context): lambda_event = {} lambda_handler(lambda_event, get_mock_context()) - wrapper.dd_tracing_enabled = True self.mock_set_correlation_ids.assert_called() self.mock_inject_correlation_ids.assert_called() @@ -453,11 +451,8 @@ def lambda_handler(event, context): ] ) + @patch("datadog_lambda.config.Config.enhanced_metrics_enabled", False) def test_no_enhanced_metrics_without_env_var(self): - patcher = patch("datadog_lambda.metric.enhanced_metrics_enabled", False) - patcher.start() - self.addCleanup(patcher.stop) - @wrapper.datadog_lambda_wrapper def lambda_handler(event, context): raise RuntimeError() @@ -565,8 +560,8 @@ def return_type_test(event, context): class TestLambdaDecoratorSettings(unittest.TestCase): + @patch("datadog_lambda.config.Config.trace_enabled", False) def test_some_envs_should_depend_on_dd_tracing_enabled(self): - wrapper.dd_tracing_enabled = False os.environ[wrapper.DD_TRACE_MANAGED_SERVICES] = "true" os.environ[wrapper.DD_ENCODE_AUTHORIZER_CONTEXT] = "true" os.environ[wrapper.DD_DECODE_AUTHORIZER_CONTEXT] = "true"