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

Author: majanjua-amzn

tests/test_lambda_context.py

@@ -85,6 +85,29 @@ 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)
+
+    assert temp_context.get_trace_entity() == dummy_segment
+
+    # 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()