[Lambda] Add unit test for lambda passthrough trace header case

Author: majanjua-amzn

tests/test_lambda_context.py

@@ -85,6 +85,31 @@ 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 compare trace ID
+    resulting_entity = temp_context.get_trace_entity()
+    assert TRACE_ID == getattr(resulting_entity, 'traceid', None)
+
+    # 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()