Added tests for testing correct sampling_context given to start_transaction by AWS integration.

Author: antonpirker

tests/integrations/aws_lambda/lambda_functions_with_embedded_sdk/TracesSampler/.gitignore

@@ -0,0 +1,11 @@
+# Need to add some ignore rules in this directory, because the unit tests will add the Sentry SDK and its dependencies 
+# into this directory to create a Lambda function package that contains everything needed to instrument a Lambda function using Sentry.
+
+# Ignore everything
+*
+
+# But not index.py
+!index.py
+
+# And not .gitignore itself
+!.gitignore

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

@@ -0,0 +1,49 @@
+import json
+import os
+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):
+    # 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"),
+    traces_sample_rate=1.0,
+    traces_sampler=trace_sampler,
+    integrations=[AwsLambdaIntegration()],
+)
+
+
+def handler(event, context):
+    # Return the sampling context data for verification
+    return {
+        "statusCode": 200,
+        "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

@@ -450,3 +450,22 @@ def test_span_origin(lambda_client, test_environment):
     assert (
         transaction_event["contexts"]["trace"]["origin"] == "auto.function.aws_lambda"
     )
+
+
+def test_traces_sampler_has_correct_sampling_context(lambda_client, test_environment):
+    """
+    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(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

@@ -1,6 +1,8 @@
 import gzip
 import json
 import os
+import shutil
+import subprocess
 import requests
 import sys
 import time
@@ -20,6 +22,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_TIMEOUT = 10
 SAM_PORT = 3001
 
@@ -47,6 +52,20 @@ def get_host_ip():
     return host
 
 
+def get_project_root():
+    """
+    Returns the absolute path to the project root directory.
+    """
+    # 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
+
+
 class LocalLambdaStack(Stack):
     """
     Uses the AWS CDK to create a local SAM stack containing Lambda functions.
@@ -124,6 +143,73 @@ def __init__(self, scope: Construct, construct_id: str, **kwargs) -> None:
                 )
             )
 
+        print(
+            "[LocalLambdaStack] Add all Lambda functions defined in "
+            "/tests/integrations/aws_lambda/lambda_functions_with_embedded_sdk/ to the SAM stack"
+        )
+        lambda_dirs = [
+            d
+            for d in os.listdir(LAMBDA_FUNCTION_WITH_EMBEDDED_SDK_DIR)
+            if os.path.isdir(os.path.join(LAMBDA_FUNCTION_WITH_EMBEDDED_SDK_DIR, d))
+        ]
+        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"
+            )
+            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 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,
+                ]
+            )
+
+            CfnResource(
+                self,
+                lambda_dir,
+                type="AWS::Serverless::Function",
+                properties={
+                    "CodeUri": os.path.join(
+                        LAMBDA_FUNCTION_WITH_EMBEDDED_SDK_DIR, lambda_dir
+                    ),
+                    "Handler": "index.handler",
+                    "Runtime": PYTHON_VERSION,
+                    "Timeout": LAMBDA_FUNCTION_TIMEOUT,
+                    "Environment": {
+                        "Variables": {
+                            "SENTRY_DSN": dsn,
+                        }
+                    },
+                },
+            )
+            print(
+                "[LocalLambdaStack] - Created Lambda function: %s (%s)"
+                % (
+                    lambda_dir,
+                    os.path.join(LAMBDA_FUNCTION_DIR, lambda_dir),
+                )
+            )
+
     @classmethod
     def wait_for_stack(cls, timeout=30, port=SAM_PORT):
         """