From d174f8d1ffebcf1e9775a22e310735bb4b59f6cf Mon Sep 17 00:00:00 2001 From: Mahad Janjua Date: Wed, 8 May 2024 10:20:39 -0700 Subject: [PATCH] [Lambda] Add unit test for lambda passthrough trace header case --- tests/test_lambda_context.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/tests/test_lambda_context.py b/tests/test_lambda_context.py index 09e08873..b7299c72 100644 --- a/tests/test_lambda_context.py +++ b/tests/test_lambda_context.py @@ -3,6 +3,7 @@ from aws_xray_sdk import global_sdk_config import pytest from aws_xray_sdk.core import lambda_launcher +from aws_xray_sdk.core.models.dummy_entities import DummySegment from aws_xray_sdk.core.models.subsegment import Subsegment @@ -85,6 +86,30 @@ def test_non_initialized(): assert temp_context.get_trace_entity() == subsegment +def test_lambda_passthrough(): + # Hold previous environment value + temp_header_var = os.environ[lambda_launcher.LAMBDA_TRACE_HEADER_KEY] + del os.environ[lambda_launcher.LAMBDA_TRACE_HEADER_KEY] + + # Set header to lambda passthrough style header + os.environ[lambda_launcher.LAMBDA_TRACE_HEADER_KEY] = "Root=%s;Lineage=10:1234abcd:3" % TRACE_ID + + temp_context = lambda_launcher.LambdaContext() + dummy_segment = temp_context.get_trace_entity() + subsegment = Subsegment("TestSubsegment", "local", dummy_segment) + temp_context.put_subsegment(subsegment) + + # Resulting entity is not the same dummy segment, so simply check that it is a dummy segment + assert isinstance(temp_context.get_trace_entity(), DummySegment) + + # Reset header value and ensure behaviour returns to normal + del os.environ[lambda_launcher.LAMBDA_TRACE_HEADER_KEY] + os.environ[lambda_launcher.LAMBDA_TRACE_HEADER_KEY] = temp_header_var + temp_context.put_subsegment(subsegment) + + assert temp_context.get_trace_entity() == subsegment + + def test_set_trace_entity(): segment = context.get_trace_entity()