Skip to content

Commit

Permalink
[Lambda] Add unit test for lambda passthrough trace header case
Browse files Browse the repository at this point in the history
  • Loading branch information
majanjua-amzn committed May 8, 2024
1 parent 3be0c9b commit d174f8d
Showing 1 changed file with 25 additions and 0 deletions.
25 changes: 25 additions & 0 deletions tests/test_lambda_context.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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()
Expand Down

0 comments on commit d174f8d

Please sign in to comment.