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..a40e2fcb 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.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..7a08d8a7 --- /dev/null +++ b/datadog_lambda/config.py @@ -0,0 +1,145 @@ +# 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 + +logger = logging.getLogger(__name__) + + +def _get_env(key, default=None, cast=None, depends_on_tracing=False): + @property + def _getter(self): + if not hasattr(self, prop_key): + val = self._resolve_env(key, default, cast, depends_on_tracing) + 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" + + +def as_list(val): + return [val.strip() for val in val.split(",") if val.strip()] + + +class Config: + def _resolve_env(self, key, default=None, cast=None, depends_on_tracing=False): + if depends_on_tracing and not self.trace_enabled: + return False + val = os.environ.get(key, default) + if cast is not None: + try: + val = cast(val) + except (ValueError, TypeError): + msg = ( + "Failed to cast environment variable '%s' with " + "value '%s' to type %s. Using default value '%s'." + ) + logger.warning(msg, key, val, cast.__name__, default) + val = default + return val + + service = _get_env("DD_SERVICE") + env = _get_env("DD_ENV") + + cold_start_tracing = _get_env( + "DD_COLD_START_TRACING", "true", as_bool, depends_on_tracing=True + ) + min_cold_start_trace_duration = _get_env("DD_MIN_COLD_START_DURATION", 3, int) + cold_start_trace_skip_lib = _get_env( + "DD_COLD_START_TRACE_SKIP_LIB", + "ddtrace.internal.compat,ddtrace.filters", + as_list, + ) + + capture_payload_max_depth = _get_env("DD_CAPTURE_LAMBDA_PAYLOAD_MAX_DEPTH", 10, int) + capture_payload_enabled = _get_env("DD_CAPTURE_LAMBDA_PAYLOAD", "false", as_bool) + + trace_enabled = _get_env("DD_TRACE_ENABLED", "true", as_bool) + make_inferred_span = _get_env( + "DD_TRACE_MANAGED_SERVICES", "true", as_bool, depends_on_tracing=True + ) + encode_authorizer_context = _get_env( + "DD_ENCODE_AUTHORIZER_CONTEXT", "true", as_bool, depends_on_tracing=True + ) + decode_authorizer_context = _get_env( + "DD_DECODE_AUTHORIZER_CONTEXT", "true", as_bool, depends_on_tracing=True + ) + add_span_pointers = _get_env("DD_BOTOCORE_ADD_SPAN_POINTERS", "true", as_bool) + trace_extractor = _get_env("DD_TRACE_EXTRACTOR") + + 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) + logs_injection = _get_env("DD_LOGS_INJECTION", "true", as_bool) + merge_xray_traces = _get_env("DD_MERGE_XRAY_TRACES", "false", as_bool) + + telemetry_enabled = _get_env( + "DD_INSTRUMENTATION_TELEMETRY_ENABLED", + "false", + as_bool, + depends_on_tracing=True, + ) + otel_enabled = _get_env("DD_TRACE_OTEL_ENABLED", "false", as_bool) + profiling_enabled = _get_env("DD_PROFILING_ENABLED", "false", as_bool) + llmobs_enabled = _get_env("DD_LLMOBS_ENABLED", "false", as_bool) + exception_replay_enabled = _get_env("DD_EXCEPTION_REPLAY_ENABLED", "false", as_bool) + data_streams_enabled = _get_env( + "DD_DATA_STREAMS_ENABLED", "false", as_bool, depends_on_tracing=True + ) + + is_gov_region = _get_env("AWS_REGION", "", lambda x: x.startswith("us-gov-")) + + local_test = _get_env("DD_LOCAL_TEST", "false", as_bool) + integration_test = _get_env("DD_INTEGRATION_TEST", "false", as_bool) + + aws_lambda_function_name = _get_env("AWS_LAMBDA_FUNCTION_NAME") + + @property + def function_name(self): + if not hasattr(self, "_config_function_name"): + if self.aws_lambda_function_name is None: + self._config_function_name = "function" + else: + self._config_function_name = self.aws_lambda_function_name + return self._config_function_name + + @property + def is_lambda_context(self): + if not hasattr(self, "_config_is_lambda_context"): + self._config_is_lambda_context = bool(self.aws_lambda_function_name) + return self._config_is_lambda_context + + @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.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..6d2af0dc 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.config 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.integration_test: 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/tag_object.py b/datadog_lambda/tag_object.py index 6d82f83b..744e4893 100644 --- a/datadog_lambda/tag_object.py +++ b/datadog_lambda/tag_object.py @@ -4,18 +4,17 @@ # Copyright 2021 Datadog, Inc. from decimal import Decimal -import logging import ujson as json +from datadog_lambda.config import config + redactable_keys = ["authorization", "x-authorization", "password", "token"] -max_depth = 10 -logger = logging.getLogger(__name__) def tag_object(span, key, obj, depth=0): if obj is None: return span.set_tag(key, obj) - if depth >= max_depth: + if depth >= config.capture_payload_max_depth: return span.set_tag(key, _redact_val(key, str(obj)[0:5000])) depth += 1 if _should_try_string(obj): diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 4b6f300a..3d5f671e 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.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 @@ -1480,7 +1470,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 0e23b721..87063411 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -18,6 +18,7 @@ is_new_sandbox, ColdStartTracer, ) +from datadog_lambda.config import config from datadog_lambda.constants import ( TraceContextSource, XraySubsegment, @@ -26,11 +27,11 @@ from datadog_lambda.module_name import modify_module_name from datadog_lambda.patch import patch_all from datadog_lambda.span_pointers import calculate_span_pointers +from datadog_lambda.tag_object import tag_object from datadog_lambda.tracing import ( 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, @@ -46,65 +47,20 @@ extract_http_status_code_tag, ) -profiling_env_var = os.environ.get("DD_PROFILING_ENABLED", "false").lower() == "true" -if profiling_env_var: +if config.profiling_enabled: from ddtrace.profiling import profiler -llmobs_env_var = os.environ.get("DD_LLMOBS_ENABLED", "false").lower() in ("true", "1") -if llmobs_env_var: +if config.llmobs_enabled: from ddtrace.llmobs import LLMObs -exception_replay_env_var = os.environ.get( - "DD_EXCEPTION_REPLAY_ENABLED", "false" -).lower() in ("true", "1") -if exception_replay_env_var: +if config.exception_replay_enabled: from ddtrace.debugging._exception.replay import SpanExceptionHandler from ddtrace.debugging._uploader import LogsIntakeUploaderV1 logger = logging.getLogger(__name__) -DD_FLUSH_TO_LOG = "DD_FLUSH_TO_LOG" -DD_LOGS_INJECTION = "DD_LOGS_INJECTION" -DD_MERGE_XRAY_TRACES = "DD_MERGE_XRAY_TRACES" -AWS_LAMBDA_FUNCTION_NAME = "AWS_LAMBDA_FUNCTION_NAME" -DD_LOCAL_TEST = "DD_LOCAL_TEST" -DD_TRACE_EXTRACTOR = "DD_TRACE_EXTRACTOR" -DD_TRACE_MANAGED_SERVICES = "DD_TRACE_MANAGED_SERVICES" -DD_ENCODE_AUTHORIZER_CONTEXT = "DD_ENCODE_AUTHORIZER_CONTEXT" -DD_DECODE_AUTHORIZER_CONTEXT = "DD_DECODE_AUTHORIZER_CONTEXT" -DD_COLD_START_TRACING = "DD_COLD_START_TRACING" -DD_MIN_COLD_START_DURATION = "DD_MIN_COLD_START_DURATION" -DD_COLD_START_TRACE_SKIP_LIB = "DD_COLD_START_TRACE_SKIP_LIB" -DD_CAPTURE_LAMBDA_PAYLOAD = "DD_CAPTURE_LAMBDA_PAYLOAD" -DD_CAPTURE_LAMBDA_PAYLOAD_MAX_DEPTH = "DD_CAPTURE_LAMBDA_PAYLOAD_MAX_DEPTH" DD_REQUESTS_SERVICE_NAME = "DD_REQUESTS_SERVICE_NAME" DD_SERVICE = "DD_SERVICE" -DD_ENV = "DD_ENV" -DD_DATA_STREAMS_ENABLED = "DD_DATA_STREAMS_ENABLED" - - -def get_env_as_int(env_key, default_value: int) -> int: - try: - return int(os.environ.get(env_key, default_value)) - except Exception as e: - logger.warn( - f"Failed to parse {env_key} as int. Using default value: {default_value}. Error: {e}" - ) - return default_value - - -dd_capture_lambda_payload_enabled = ( - os.environ.get(DD_CAPTURE_LAMBDA_PAYLOAD, "false").lower() == "true" -) - -if dd_capture_lambda_payload_enabled: - import datadog_lambda.tag_object as tag_object - - tag_object.max_depth = get_env_as_int( - DD_CAPTURE_LAMBDA_PAYLOAD_MAX_DEPTH, tag_object.max_depth - ) - -env_env_var = os.environ.get(DD_ENV, None) init_timestamp_ns = time_ns() @@ -161,59 +117,16 @@ def __init__(self, func): """Executes when the wrapped function gets wrapped""" try: self.func = func - self.flush_to_log = os.environ.get(DD_FLUSH_TO_LOG, "").lower() == "true" - self.logs_injection = ( - os.environ.get(DD_LOGS_INJECTION, "true").lower() == "true" - ) - self.merge_xray_traces = ( - os.environ.get(DD_MERGE_XRAY_TRACES, "false").lower() == "true" - ) - self.function_name = os.environ.get(AWS_LAMBDA_FUNCTION_NAME, "function") - self.service = os.environ.get(DD_SERVICE, None) - self.extractor_env = os.environ.get(DD_TRACE_EXTRACTOR, None) self.trace_extractor = None self.span = None self.inferred_span = None - depends_on_dd_tracing_enabled = ( - lambda original_boolean: dd_tracing_enabled and original_boolean - ) - self.make_inferred_span = depends_on_dd_tracing_enabled( - os.environ.get(DD_TRACE_MANAGED_SERVICES, "true").lower() == "true" - ) - self.encode_authorizer_context = depends_on_dd_tracing_enabled( - os.environ.get(DD_ENCODE_AUTHORIZER_CONTEXT, "true").lower() == "true" - ) - self.decode_authorizer_context = depends_on_dd_tracing_enabled( - os.environ.get(DD_DECODE_AUTHORIZER_CONTEXT, "true").lower() == "true" - ) - self.cold_start_tracing = depends_on_dd_tracing_enabled( - os.environ.get(DD_COLD_START_TRACING, "true").lower() == "true" - ) - self.min_cold_start_trace_duration = get_env_as_int( - DD_MIN_COLD_START_DURATION, 3 - ) - self.data_streams_enabled = ( - os.environ.get(DD_DATA_STREAMS_ENABLED, "false").lower() == "true" - ) - self.local_testing_mode = os.environ.get( - DD_LOCAL_TEST, "false" - ).lower() in ("true", "1") - self.cold_start_trace_skip_lib = [ - "ddtrace.internal.compat", - "ddtrace.filters", - ] - if DD_COLD_START_TRACE_SKIP_LIB in os.environ: - try: - self.cold_start_trace_skip_lib = os.environ[ - DD_COLD_START_TRACE_SKIP_LIB - ].split(",") - except Exception: - logger.debug(f"Malformatted for env {DD_COLD_START_TRACE_SKIP_LIB}") self.response = None - if profiling_env_var: - self.prof = profiler.Profiler(env=env_env_var, service=self.service) - if self.extractor_env: - extractor_parts = self.extractor_env.rsplit(".", 1) + + if config.profiling_enabled: + self.prof = profiler.Profiler(env=config.env, service=config.service) + + if config.trace_extractor: + extractor_parts = config.trace_extractor.rsplit(".", 1) if len(extractor_parts) == 2: (mod_name, extractor_name) = extractor_parts modified_extractor_name = modify_module_name(mod_name) @@ -221,7 +134,7 @@ def __init__(self, func): self.trace_extractor = getattr(extractor_module, extractor_name) # Inject trace correlation ids to logs - if self.logs_injection: + if config.logs_injection: inject_correlation_ids() # This prevents a breaking change in ddtrace v0.49 regarding the service name @@ -233,11 +146,11 @@ def __init__(self, func): patch_all() # Enable LLM Observability - if llmobs_env_var: + if config.llmobs_enabled: LLMObs.enable() # Enable Exception Replay - if exception_replay_env_var: + if config.exception_replay_enabled: logger.debug("Enabling exception replay") SpanExceptionHandler.enable() @@ -307,7 +220,7 @@ def _before(self, event, context): event, context, extractor=self.trace_extractor, - decode_authorizer_context=self.decode_authorizer_context, + decode_authorizer_context=config.decode_authorizer_context, ) self.event_source = event_source # Create a Datadog X-Ray subsegment with the trace context @@ -321,28 +234,28 @@ def _before(self, event, context): XraySubsegment.TRACE_KEY, ) - if dd_tracing_enabled: - set_dd_trace_py_root(trace_context_source, self.merge_xray_traces) - if self.make_inferred_span: + if config.trace_enabled: + set_dd_trace_py_root(trace_context_source, config.merge_xray_traces) + if config.make_inferred_span: self.inferred_span = create_inferred_span( - event, context, event_source, self.decode_authorizer_context + event, context, event_source, config.decode_authorizer_context ) - if self.data_streams_enabled: + if config.data_streams_enabled: set_dsm_context(event, event_source) self.span = create_function_execution_span( context=context, - function_name=self.function_name, + function_name=config.function_name, is_cold_start=is_cold_start(), is_proactive_init=is_proactive_init(), trace_context_source=trace_context_source, - merge_xray_traces=self.merge_xray_traces, + merge_xray_traces=config.merge_xray_traces, trigger_tags=self.trigger_tags, parent_span=self.inferred_span, span_pointers=calculate_span_pointers(event_source, event), ) else: set_correlation_ids() - if profiling_env_var and is_new_sandbox(): + if config.profiling_enabled and is_new_sandbox(): self.prof.start(stop_on_exit=False, profile_children=True) logger.debug("datadog_lambda_wrapper _before() done") except Exception as e: @@ -361,14 +274,14 @@ def _after(self, event, context): create_dd_dummy_metadata_subsegment( self.trigger_tags, XraySubsegment.LAMBDA_FUNCTION_TAGS_KEY ) - should_trace_cold_start = self.cold_start_tracing and is_new_sandbox() + should_trace_cold_start = config.cold_start_tracing and is_new_sandbox() if should_trace_cold_start: trace_ctx = tracer.current_trace_context() if self.span: - if dd_capture_lambda_payload_enabled: - tag_object.tag_object(self.span, "function.request", event) - tag_object.tag_object(self.span, "function.response", self.response) + if config.capture_payload_enabled: + tag_object(self.span, "function.request", event) + tag_object(self.span, "function.response", self.response) if status_code: self.span.set_tag("http.status_code", status_code) @@ -378,8 +291,8 @@ def _after(self, event, context): if status_code: self.inferred_span.set_tag("http.status_code", status_code) - if self.service: - self.inferred_span.set_tag("peer.service", self.service) + if config.service: + self.inferred_span.set_tag("peer.service", config.service) if InferredSpanInfo.is_async(self.inferred_span) and self.span: self.inferred_span.finish(finish_time=self.span.start) @@ -391,33 +304,35 @@ def _after(self, event, context): following_span = self.span or self.inferred_span ColdStartTracer( tracer, - self.function_name, + config.function_name, following_span.start_ns, trace_ctx, - self.min_cold_start_trace_duration, - self.cold_start_trace_skip_lib, + config.min_cold_start_trace_duration, + config.cold_start_trace_skip_lib, ).trace() except Exception as e: logger.debug("Failed to create cold start spans. %s", e) - if not self.flush_to_log or should_use_extension: + if not config.flush_to_log or should_use_extension: from datadog_lambda.metric import flush_stats flush_stats(context) - if should_use_extension and self.local_testing_mode: + if should_use_extension and config.local_test: # when testing locally, the extension does not know when an # invocation completes because it does not have access to the # logs api flush_extension() - if llmobs_env_var: + if config.llmobs_enabled: LLMObs.flush() # Flush exception replay - if exception_replay_env_var: + if config.exception_replay_enabled: LogsIntakeUploaderV1._instance.periodic() - if self.encode_authorizer_context and is_authorizer_response(self.response): + if config.encode_authorizer_context and is_authorizer_response( + self.response + ): self._inject_authorizer_span_headers( event.get("requestContext", {}).get("requestId") ) diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 00000000..33869802 --- /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..d75b5f43 100644 --- a/tests/test_cold_start.py +++ b/tests/test_cold_start.py @@ -8,6 +8,8 @@ import datadog_lambda.cold_start as cold_start import datadog_lambda.wrapper as wrapper +from tests.utils import get_mock_context + class TestColdStartTracingSetup(unittest.TestCase): def test_proactive_init(self): @@ -247,7 +249,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" ) @@ -257,10 +259,7 @@ def finish(span): def handler(event, context): import tabnanny - lambda_context = MagicMock() - lambda_context.invoked_function_arn = ( - "arn:aws:lambda:us-west-1:123457598159:function:python-layer-test:1" - ) + lambda_context = get_mock_context() handler.cold_start_tracing = True handler({}, lambda_context) diff --git a/tests/test_config.py b/tests/test_config.py new file mode 100644 index 00000000..92002439 --- /dev/null +++ b/tests/test_config.py @@ -0,0 +1,240 @@ +import pytest + +from datadog_lambda.config import config, _get_env, 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 + + +def _test_as_bool(env_key, conf_key, default): + return ( + (env_key, conf_key, None, default), + (env_key, conf_key, "", False), + (env_key, conf_key, "true", True), + (env_key, conf_key, "TRUE", True), + (env_key, conf_key, "false", False), + (env_key, conf_key, "FALSE", False), + (env_key, conf_key, "1", True), + (env_key, conf_key, "0", False), + (env_key, conf_key, "purple", False), + ) + + +def _test_int(env_key, conf_key, default): + return ( + (env_key, conf_key, None, default), + (env_key, conf_key, "", default), + (env_key, conf_key, "5", 5), + (env_key, conf_key, "0", 0), + (env_key, conf_key, "2.5", default), + (env_key, conf_key, "-1", -1), + (env_key, conf_key, "purple", default), + ) + + +def _test_as_list(env_key, conf_key, default): + return ( + (env_key, conf_key, None, default.split(",")), + (env_key, conf_key, "", []), + (env_key, conf_key, " ", []), + (env_key, conf_key, ",", []), + (env_key, conf_key, " , ", []), + (env_key, conf_key, "a", ["a"]), + (env_key, conf_key, "a,", ["a"]), + (env_key, conf_key, "a, ", ["a"]), + (env_key, conf_key, "a,b", ["a", "b"]), + (env_key, conf_key, "a, b", ["a", "b"]), + ) + + +_test_config_from_environ = ( + *_test_as_bool("DD_FLUSH_TO_LOG", "flush_to_log", default=False), + *_test_as_bool("DD_LOGS_INJECTION", "logs_injection", default=True), + *_test_as_bool("DD_TRACE_ENABLED", "trace_enabled", default=True), + *_test_as_bool("DD_COLD_START_TRACING", "cold_start_tracing", default=True), + *_test_as_bool("DD_TRACE_MANAGED_SERVICES", "make_inferred_span", default=True), + *_test_as_bool( + "DD_ENCODE_AUTHORIZER_CONTEXT", "encode_authorizer_context", default=True + ), + *_test_as_bool( + "DD_DECODE_AUTHORIZER_CONTEXT", "decode_authorizer_context", default=True + ), + *_test_as_bool("DD_FLUSH_IN_THREAD", "flush_in_thread", default=False), + *_test_as_bool("DD_ENHANCED_METRICS", "enhanced_metrics_enabled", default=True), + *_test_as_bool("DD_INTEGRATION_TEST", "integration_test", default=False), + *_test_as_bool("DD_BOTOCORE_ADD_SPAN_POINTERS", "add_span_pointers", default=True), + *_test_as_bool("DD_TRACE_OTEL_ENABLED", "otel_enabled", default=False), + *_test_as_bool( + "DD_INSTRUMENTATION_TELEMETRY_ENABLED", "telemetry_enabled", default=False + ), + *_test_as_bool("DD_MERGE_XRAY_TRACES", "merge_xray_traces", default=False), + *_test_as_bool("DD_PROFILING_ENABLED", "profiling_enabled", default=False), + *_test_as_bool("DD_LLMOBS_ENABLED", "llmobs_enabled", default=False), + *_test_as_bool( + "DD_EXCEPTION_REPLAY_ENABLED", "exception_replay_enabled", default=False + ), + *_test_as_bool( + "DD_CAPTURE_LAMBDA_PAYLOAD", "capture_payload_enabled", default=False + ), + *_test_as_bool("DD_LOCAL_TEST", "local_test", default=False), + *_test_as_bool("DD_DATA_STREAMS_ENABLED", "data_streams_enabled", default=False), + *_test_int( + "DD_CAPTURE_LAMBDA_PAYLOAD_MAX_DEPTH", "capture_payload_max_depth", default=10 + ), + *_test_int( + "DD_MIN_COLD_START_DURATION", "min_cold_start_trace_duration", default=3 + ), + *_test_as_list( + "DD_COLD_START_TRACE_SKIP_LIB", + "cold_start_trace_skip_lib", + default="ddtrace.internal.compat,ddtrace.filters", + ), + ("DD_SERVICE", "service", None, None), + ("DD_SERVICE", "service", "", ""), + ("DD_SERVICE", "service", "my_service", "my_service"), + ("AWS_LAMBDA_FUNCTION_NAME", "aws_lambda_function_name", None, None), + ("AWS_LAMBDA_FUNCTION_NAME", "aws_lambda_function_name", "", ""), + ( + "AWS_LAMBDA_FUNCTION_NAME", + "aws_lambda_function_name", + "my_function", + "my_function", + ), + ("AWS_LAMBDA_FUNCTION_NAME", "function_name", None, "function"), + ("AWS_LAMBDA_FUNCTION_NAME", "function_name", "", ""), + ("AWS_LAMBDA_FUNCTION_NAME", "function_name", "my_function", "my_function"), + ("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), + ("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_TRACE_EXTRACTOR", "trace_extractor", None, None), + ("DD_TRACE_EXTRACTOR", "trace_extractor", "", ""), + ("DD_TRACE_EXTRACTOR", "trace_extractor", "my_extractor", "my_extractor"), + ("DD_ENV", "env", None, None), + ("DD_ENV", "env", "", ""), + ("DD_ENV", "env", "my_env", "my_env"), +) + + +@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_config_from_environ_depends_on_tracing = ( + *_test_as_bool("DD_COLD_START_TRACING", "cold_start_tracing", default=True), + *_test_as_bool("DD_TRACE_MANAGED_SERVICES", "make_inferred_span", default=True), + *_test_as_bool( + "DD_ENCODE_AUTHORIZER_CONTEXT", "encode_authorizer_context", default=True + ), + *_test_as_bool( + "DD_DECODE_AUTHORIZER_CONTEXT", "decode_authorizer_context", default=True + ), + *_test_as_bool("DD_DATA_STREAMS_ENABLED", "data_streams_enabled", default=False), + *_test_as_bool( + "DD_INSTRUMENTATION_TELEMETRY_ENABLED", "telemetry_enabled", default=False + ), +) + + +@pytest.mark.parametrize( + "env_key,conf_key,env_val,conf_val", _test_config_from_environ_depends_on_tracing +) +def test_config_from_environ_depends_on_tracing( + env_key, conf_key, env_val, conf_val, setenv +): + setenv(env_key, env_val) + setenv("DD_TRACE_ENABLED", "false") + assert getattr(config, conf_key) is False + + +def test_config_aws_lambda_function_name(setenv): + # these config values all access the same environment variable, test to + # ensure the wrong value is not cached + setenv("AWS_LAMBDA_FUNCTION_NAME", "my_function") + assert config.aws_lambda_function_name == "my_function" + assert config.function_name == "my_function" + assert config.is_lambda_context is True + + +_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 + + +def test__get_env_does_not_log_when_env_not_set(setenv, monkeypatch): + setenv("TEST_1", None) + setenv("TEST_2", None) + setenv("TEST_3", None) + setenv("TEST_4", None) + + class Testing(Config): + test_1 = _get_env("TEST_1") + test_2 = _get_env("TEST_2", "purple") + test_3 = _get_env("TEST_3", "true", bool) + test_4 = _get_env("TEST_4", "true", bool, depends_on_tracing=True) + + logs = [] + + def cap_warn(*args, **kwargs): + logs.append(args) + + monkeypatch.setattr("datadog_lambda.config.logger.warning", cap_warn) + + testing = Testing() + testing.test_1 + testing.test_2 + testing.test_3 + testing.test_4 + + assert not logs diff --git a/tests/test_metric.py b/tests/test_metric.py index e7dab2c3..aa537d34 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.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_patch.py b/tests/test_patch.py index bf924875..b03d2e23 100644 --- a/tests/test_patch.py +++ b/tests/test_patch.py @@ -1,3 +1,4 @@ +import pytest import unittest from unittest.mock import patch, MagicMock @@ -5,6 +6,13 @@ from datadog_lambda.patch import _patch_http, _ensure_patch_requests from datadog_lambda.constants import TraceHeader +from ddtrace.contrib.internal.requests.patch import unpatch as unpatch_requests + + +@pytest.fixture(scope="module", autouse=True) +def reset_patches(): + unpatch_requests() + class TestPatchHTTPClients(unittest.TestCase): def setUp(self): diff --git a/tests/test_tag_object.py b/tests/test_tag_object.py index 77512164..574bb331 100644 --- a/tests/test_tag_object.py +++ b/tests/test_tag_object.py @@ -29,6 +29,7 @@ def test_tag_object(self): True, ) + @patch("datadog_lambda.config.Config.capture_payload_max_depth", 2) def test_tag_object_max_depth(self): payload = { "hello": "world", @@ -41,11 +42,8 @@ def test_tag_object_max_depth(self): "vals": [{"thingOne": 1}, {"thingTwo": 2}], } spanMock = MagicMock() - import datadog_lambda.tag_object as lib_ref - lib_ref.max_depth = 2 # setting up the test tag_object(spanMock, "function.request", payload) - lib_ref.max_depth = 10 # revert the setup spanMock.set_tag.assert_has_calls( [ call("function.request.vals.0", "{'thingOne': 1}"), @@ -62,6 +60,7 @@ def test_tag_object_max_depth(self): True, ) + @patch("datadog_lambda.config.Config.capture_payload_max_depth", 0) def test_tag_object_max_depth_0(self): payload = { "hello": "world", @@ -74,11 +73,8 @@ def test_tag_object_max_depth_0(self): "vals": [{"thingOne": 1}, {"thingTwo": 2}], } spanMock = MagicMock() - import datadog_lambda.tag_object as lib_ref - lib_ref.max_depth = 0 # setting up the test tag_object(spanMock, "function.request", payload) - lib_ref.max_depth = 10 # revert the setup spanMock.set_tag.assert_has_calls( [ call( diff --git a/tests/test_tracing.py b/tests/test_tracing.py index e38e4ecd..a629343e 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -251,20 +251,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") @@ -984,11 +980,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 +1121,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 +1138,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 +1238,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 +1279,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 +2382,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 +2423,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 f482fa3d..f0240905 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" ) @@ -80,9 +81,8 @@ def setUp(self): self.mock_set_dsm_context = 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) @@ -92,7 +92,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( @@ -189,9 +188,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): @@ -199,7 +198,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() @@ -457,11 +455,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() @@ -515,6 +510,7 @@ def lambda_handler(event, context): self.assertEqual(os.environ.get("DD_REQUESTS_SERVICE_NAME"), "myAwesomeService") del os.environ["DD_SERVICE"] + @patch("datadog_lambda.config.Config.make_inferred_span", False) def test_encode_authorizer_span(self): @wrapper.datadog_lambda_wrapper def lambda_handler(event, context): @@ -541,7 +537,6 @@ def lambda_handler(event, context): trace_ctx.sampling_priority = 1 test_span.finish() lambda_handler.inferred_span = test_span - lambda_handler.make_inferred_span = False result = lambda_handler(lambda_event, lambda_context) raw_inject_data = result["context"]["_datadog"] self.assertIsInstance(raw_inject_data, str) @@ -569,7 +564,7 @@ def return_type_test(event, context): def test_set_dsm_context_called_when_DSM_and_tracing_enabled(self): os.environ["DD_DATA_STREAMS_ENABLED"] = "true" - wrapper.dd_tracing_enabled = True + os.environ["DD_TRACE_ENABLED"] = "true" @wrapper.datadog_lambda_wrapper def lambda_handler(event, context): @@ -583,7 +578,7 @@ def lambda_handler(event, context): def test_set_dsm_context_not_called_when_only_DSM_enabled(self): os.environ["DD_DATA_STREAMS_ENABLED"] = "true" - wrapper.dd_tracing_enabled = False + os.environ["DD_TRACE_ENABLED"] = "false" @wrapper.datadog_lambda_wrapper def lambda_handler(event, context): @@ -597,7 +592,7 @@ def lambda_handler(event, context): def test_set_dsm_context_not_called_when_only_tracing_enabled(self): os.environ["DD_DATA_STREAMS_ENABLED"] = "false" - wrapper.dd_tracing_enabled = True + os.environ["DD_TRACE_ENABLED"] = "true" @wrapper.datadog_lambda_wrapper def lambda_handler(event, context): @@ -611,7 +606,7 @@ def lambda_handler(event, context): def test_set_dsm_context_not_called_when_tracing_and_DSM_disabled(self): os.environ["DD_DATA_STREAMS_ENABLED"] = "false" - wrapper.dd_tracing_enabled = False + os.environ["DD_TRACE_ENABLED"] = "false" @wrapper.datadog_lambda_wrapper def lambda_handler(event, context): @@ -624,18 +619,6 @@ def lambda_handler(event, context): del os.environ["DD_DATA_STREAMS_ENABLED"] -class TestLambdaDecoratorSettings(unittest.TestCase): - 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" - decorator = wrapper._LambdaDecorator(func=None) - self.assertFalse(decorator.make_inferred_span) - self.assertFalse(decorator.encode_authorizer_context) - self.assertFalse(decorator.decode_authorizer_context) - - class TestLambdaWrapperWithTraceContext(unittest.TestCase): xray_root = "1-5e272390-8c398be037738dc042009320" xray_parent = "94ae789b969f1cc5" @@ -706,14 +689,28 @@ def handler(event, context): class TestLambdaWrapperFlushExtension(unittest.TestCase): - def setUp(self): - self.orig_environ = os.environ + @patch("datadog_lambda.config.Config.local_test", True) + @patch("datadog_lambda.wrapper.should_use_extension", True) + def test_local_test_true_flushing(self): + flushes = [] + lambda_event = {} + lambda_context = get_mock_context() + + def flush(): + flushes.append(1) - def tearDown(self): - os.environ = self.orig_environ + @patch("datadog_lambda.wrapper.flush_extension", flush) + @wrapper.datadog_lambda_wrapper + def lambda_handler(event, context): + pass + lambda_handler(lambda_event, lambda_context) + + self.assertEqual(len(flushes), 1) + + @patch("datadog_lambda.config.Config.local_test", False) @patch("datadog_lambda.wrapper.should_use_extension", True) - def test_local_test_envvar_flushing(self): + def test_local_test_false_flushing(self): flushes = [] lambda_event = {} lambda_context = get_mock_context() @@ -721,24 +718,11 @@ def test_local_test_envvar_flushing(self): def flush(): flushes.append(1) - for environ, flush_called in ( - ({"DD_LOCAL_TEST": "True"}, True), - ({"DD_LOCAL_TEST": "true"}, True), - ({"DD_LOCAL_TEST": "1"}, True), - ({"DD_LOCAL_TEST": "False"}, False), - ({"DD_LOCAL_TEST": "false"}, False), - ({"DD_LOCAL_TEST": "0"}, False), - ({"DD_LOCAL_TEST": ""}, False), - ({}, False), - ): - os.environ = environ - flushes.clear() - - @patch("datadog_lambda.wrapper.flush_extension", flush) - @wrapper.datadog_lambda_wrapper - def lambda_handler(event, context): - pass + @patch("datadog_lambda.wrapper.flush_extension", flush) + @wrapper.datadog_lambda_wrapper + def lambda_handler(event, context): + pass - lambda_handler(lambda_event, lambda_context) + lambda_handler(lambda_event, lambda_context) - self.assertEqual(flush_called, len(flushes) == 1) + self.assertEqual(len(flushes), 0) diff --git a/tests/utils.py b/tests/utils.py index 0f246e68..2d56ca0c 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -22,6 +22,7 @@ def get_mock_context( lambda_context.invoked_function_arn = invoked_function_arn lambda_context.function_version = function_version lambda_context.function_name = function_name + lambda_context.get_remaining_time_in_millis = lambda: 100 lambda_context.client_context = ClientContext(custom) return lambda_context