From 05818f98b6050fcb5b01a23249a28da3a1832edf Mon Sep 17 00:00:00 2001 From: Michael Brewer Date: Tue, 10 Aug 2021 10:57:49 -0700 Subject: [PATCH] feat(tracer): Autodetect lambda execution env --- aws_lambda_powertools/shared/constants.py | 3 +-- aws_lambda_powertools/tracing/tracer.py | 7 +++---- tests/functional/test_tracing.py | 21 +++------------------ tests/unit/test_tracing.py | 10 +++++++--- 4 files changed, 14 insertions(+), 27 deletions(-) diff --git a/aws_lambda_powertools/shared/constants.py b/aws_lambda_powertools/shared/constants.py index 8388eded654..6061462a051 100644 --- a/aws_lambda_powertools/shared/constants.py +++ b/aws_lambda_powertools/shared/constants.py @@ -12,10 +12,9 @@ EVENT_HANDLER_DEBUG_ENV: str = "POWERTOOLS_EVENT_HANDLER_DEBUG" -SAM_LOCAL_ENV: str = "AWS_SAM_LOCAL" -CHALICE_LOCAL_ENV: str = "AWS_CHALICE_CLI_MODE" SERVICE_NAME_ENV: str = "POWERTOOLS_SERVICE_NAME" XRAY_TRACE_ID_ENV: str = "_X_AMZN_TRACE_ID" +LAMBDA_TASK_ROOT_ENV: str = "LAMBDA_TASK_ROOT" XRAY_SDK_MODULE: str = "aws_xray_sdk" XRAY_SDK_CORE_MODULE: str = "aws_xray_sdk.core" diff --git a/aws_lambda_powertools/tracing/tracer.py b/aws_lambda_powertools/tracing/tracer.py index bd4244b3171..e57d24044dc 100644 --- a/aws_lambda_powertools/tracing/tracer.py +++ b/aws_lambda_powertools/tracing/tracer.py @@ -719,16 +719,15 @@ def _is_tracer_disabled() -> Union[bool, str]: Union[bool, str] """ logger.debug("Verifying whether Tracing has been disabled") - is_lambda_sam_cli = os.getenv(constants.SAM_LOCAL_ENV) - is_chalice_cli = os.getenv(constants.CHALICE_LOCAL_ENV) + is_lambda_env = os.getenv(constants.LAMBDA_TASK_ROOT_ENV) is_disabled = resolve_truthy_env_var_choice(env=os.getenv(constants.TRACER_DISABLED_ENV, "false")) if is_disabled: logger.debug("Tracing has been disabled via env var POWERTOOLS_TRACE_DISABLED") return is_disabled - if is_lambda_sam_cli or is_chalice_cli: - logger.debug("Running under SAM CLI env or not in Lambda env; disabling Tracing") + if not is_lambda_env: + logger.debug("Running outside Lambda env; disabling Tracing") return True return False diff --git a/tests/functional/test_tracing.py b/tests/functional/test_tracing.py index 617bf816f86..b330ab6316f 100644 --- a/tests/functional/test_tracing.py +++ b/tests/functional/test_tracing.py @@ -60,26 +60,11 @@ def greeting(name, message): greeting(name="Foo", message="Bar") -def test_tracer_lambda_emulator(monkeypatch, dummy_response): - # GIVEN tracer runs locally - monkeypatch.setenv("AWS_SAM_LOCAL", "true") +def test_tracer_lambda_outside_lambda_env(monkeypatch, dummy_response): + # GIVEN tracer runs locally (ie: `LAMBDA_TASK_ROOT` is not set) tracer = Tracer() - # WHEN a lambda function is run through SAM CLI - @tracer.capture_lambda_handler - def handler(event, context): - return dummy_response - - # THEN tracer should run in disabled mode, and not raise an Exception - handler({}, {}) - - -def test_tracer_chalice_cli_mode(monkeypatch, dummy_response): - # GIVEN tracer runs locally - monkeypatch.setenv("AWS_CHALICE_CLI_MODE", "true") - tracer = Tracer() - - # WHEN a lambda function is run through the Chalice CLI. + # WHEN a lambda function is run through outside of a lambda @tracer.capture_lambda_handler def handler(event, context): return dummy_response diff --git a/tests/unit/test_tracing.py b/tests/unit/test_tracing.py index fdfdf5c6d6e..2b147ec4405 100644 --- a/tests/unit/test_tracing.py +++ b/tests/unit/test_tracing.py @@ -125,9 +125,11 @@ def greeting(name, message): ) -def test_tracer_custom_metadata(mocker, dummy_response, provider_stub): +def test_tracer_custom_metadata(monkeypatch, mocker, dummy_response, provider_stub): # GIVEN Tracer is initialized with booking as the service name + monkeypatch.setenv("LAMBDA_TASK_ROOT", "/opt/") put_metadata_mock = mocker.MagicMock() + provider = provider_stub(put_metadata_mock=put_metadata_mock) tracer = Tracer(provider=provider, service="booking") @@ -143,8 +145,9 @@ def test_tracer_custom_metadata(mocker, dummy_response, provider_stub): ) -def test_tracer_custom_annotation(mocker, dummy_response, provider_stub): +def test_tracer_custom_annotation(monkeypatch, mocker, dummy_response, provider_stub): # GIVEN Tracer is initialized + monkeypatch.setenv("LAMBDA_TASK_ROOT", "/opt/") put_annotation_mock = mocker.MagicMock() provider = provider_stub(put_annotation_mock=put_annotation_mock) tracer = Tracer(provider=provider) @@ -214,8 +217,9 @@ def greeting(name, message): @mock.patch("aws_lambda_powertools.tracing.tracer.aws_xray_sdk.core.patch") -def test_tracer_patch_modules(xray_patch_mock, mocker): +def test_tracer_patch_modules(xray_patch_mock, monkeypatch, mocker): # GIVEN tracer is initialized with a list of modules to patch + monkeypatch.setenv("LAMBDA_TASK_ROOT", "/opt/") modules = ["boto3"] # WHEN modules are supported by X-Ray