Testing for sampling_context set by aws lambda integration

Author: antonpirker

tests/integrations/aws_lambda/lambda_functions_with_embedded_sdk/TracesSampler/index.py

@@ -3,9 +3,29 @@
 import sentry_sdk
 from sentry_sdk.integrations.aws_lambda import AwsLambdaIntegration
 
+# Global variables to store sampling context for verification
+sampling_context_data = {
+    "aws_event_present": False,
+    "aws_context_present": False,
+    "event_data": None,
+}
+
+
 def trace_sampler(sampling_context):
-    print("I AM IN TRACES SAMPLER")
-    return 1
+    # Store the sampling context for verification
+    global sampling_context_data
+
+    # Check if aws_event and aws_context are in the sampling_context
+    if "aws_event" in sampling_context:
+        sampling_context_data["aws_event_present"] = True
+        sampling_context_data["event_data"] = sampling_context["aws_event"]
+
+    if "aws_context" in sampling_context:
+        sampling_context_data["aws_context_present"] = True
+
+    print("Sampling context data:", sampling_context_data)
+    return 1.0  # Always sample
+
 
 sentry_sdk.init(
     dsn=os.environ.get("SENTRY_DSN"),
@@ -14,12 +34,16 @@ def trace_sampler(sampling_context):
     integrations=[AwsLambdaIntegration()],
 )
 
+
 def handler(event, context):
-    # This is a simple Lambda function with Sentry SDK embedded
+    # Return the sampling context data for verification
     return {
         "statusCode": 200,
-        "body": json.dumps({
-            "message": "Hello from Lambda with embedded Sentry SDK!",
-            "event": event
-        })
+        "body": json.dumps(
+            {
+                "message": "Hello from Lambda with embedded Sentry SDK!",
+                "event": event,
+                "sampling_context_data": sampling_context_data,
+            }
+        ),
     }

tests/integrations/aws_lambda/test_aws_lambda.py

@@ -453,12 +453,19 @@ def test_span_origin(lambda_client, test_environment):
 
 
 def test_traces_sampler_has_correct_sampling_context(lambda_client, test_environment):
-    lambda_client.invoke(
+    """
+    Test that aws_event and aws_context are passed in the custom_sampling_context
+    when using the AWS Lambda integration.
+    """
+    test_payload = {"test_key": "test_value"}
+    response = lambda_client.invoke(
         FunctionName="TracesSampler",
-        Payload=json.dumps({}),
-    )
-    envelopes = test_environment["server"].envelopes
-    (transaction_event,) = envelopes
-    assert (
-        transaction_event["contexts"]["trace"]["origin"] == "auto.function.aws_lambda"
+        Payload=json.dumps(test_payload),
     )
+    response_payload = json.loads(response["Payload"].read().decode())
+    sampling_context_data = json.loads(response_payload["body"])[
+        "sampling_context_data"
+    ]
+    assert sampling_context_data.get("aws_event_present") is True
+    assert sampling_context_data.get("aws_context_present") is True
+    assert sampling_context_data.get("event_data", {}).get("test_key") == "test_value"

tests/integrations/aws_lambda/utils.py

@@ -23,7 +23,9 @@
 
 
 LAMBDA_FUNCTION_DIR = "./tests/integrations/aws_lambda/lambda_functions/"
-LAMBDA_FUNCTION_WITH_EMBEDDED_SDK_DIR = "./tests/integrations/aws_lambda/lambda_functions_with_embedded_sdk/"
+LAMBDA_FUNCTION_WITH_EMBEDDED_SDK_DIR = (
+    "./tests/integrations/aws_lambda/lambda_functions_with_embedded_sdk/"
+)
 LAMBDA_FUNCTION_TIMEOUT = 10
 SAM_PORT = 3001
 
@@ -57,11 +59,11 @@ def get_project_root():
     """
     # Start from the current file's directory
     current_dir = os.path.dirname(os.path.abspath(__file__))
-    
+
     # Navigate up to the project root (4 levels up from tests/integrations/aws_lambda/)
     # This is equivalent to the multiple dirname() calls
     project_root = os.path.abspath(os.path.join(current_dir, "../../../"))
-    
+
     return project_root
 
 
@@ -153,35 +155,44 @@ def __init__(self, scope: Construct, construct_id: str, **kwargs) -> None:
         ]
         for lambda_dir in lambda_dirs:
             # Copy the Sentry SDK into the function directory
-            sdk_path = os.path.join(LAMBDA_FUNCTION_WITH_EMBEDDED_SDK_DIR, lambda_dir, "sentry_sdk")
+            sdk_path = os.path.join(
+                LAMBDA_FUNCTION_WITH_EMBEDDED_SDK_DIR, lambda_dir, "sentry_sdk"
+            )
             if not os.path.exists(sdk_path):
                 # Find the Sentry SDK in the current environment
                 import sentry_sdk as sdk_module
+
                 sdk_source = os.path.dirname(sdk_module.__file__)
                 shutil.copytree(sdk_source, sdk_path)
 
-            # Install the requirements of Sentry SDK into the function directory           
-            requirements_file = os.path.join(get_project_root(), "requirements-aws-lambda-layer.txt")
-            
+            # Install the requirements of Sentry SDK into the function directory
+            requirements_file = os.path.join(
+                get_project_root(), "requirements-aws-lambda-layer.txt"
+            )
+
             # Install the package using pip
-            subprocess.check_call([
-                sys.executable, 
-                "-m", 
-                "pip", 
-                "install", 
-                "--upgrade",
-                "--target", 
-                os.path.join(LAMBDA_FUNCTION_WITH_EMBEDDED_SDK_DIR, lambda_dir),
-                "-r", 
-                requirements_file,
-            ])            
-            
+            subprocess.check_call(
+                [
+                    sys.executable,
+                    "-m",
+                    "pip",
+                    "install",
+                    "--upgrade",
+                    "--target",
+                    os.path.join(LAMBDA_FUNCTION_WITH_EMBEDDED_SDK_DIR, lambda_dir),
+                    "-r",
+                    requirements_file,
+                ]
+            )
+
             CfnResource(
                 self,
                 lambda_dir,
                 type="AWS::Serverless::Function",
                 properties={
-                    "CodeUri": os.path.join(LAMBDA_FUNCTION_WITH_EMBEDDED_SDK_DIR, lambda_dir),
+                    "CodeUri": os.path.join(
+                        LAMBDA_FUNCTION_WITH_EMBEDDED_SDK_DIR, lambda_dir
+                    ),
                     "Handler": "index.handler",
                     "Runtime": PYTHON_VERSION,
                     "Timeout": LAMBDA_FUNCTION_TIMEOUT,