Skip to content

Commit 77190ae

Browse files
maxhoheisermaxhoheiser
committed
feat: add x-ray trace id propagation event api gateway
1 parent 96ce43c commit 77190ae

File tree

1 file changed

+219
-0
lines changed

1 file changed

+219
-0
lines changed

‎tests/aws/services/events/test_x_ray_trace_propagation.py

Lines changed: 219 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,219 @@
1+
import json
2+
3+
from localstack import config
4+
from localstack.aws.api.lambda_ import Runtime
5+
from localstack.testing.pytest import markers
6+
from localstack.utils.strings import short_uid
7+
from localstack.utils.sync import retry
8+
from localstack.utils.testutil import check_expected_lambda_log_events_length
9+
10+
APIGATEWAY_ASSUME_ROLE_POLICY = {
11+
"Statement": {
12+
"Sid": "",
13+
"Effect": "Allow",
14+
"Principal": {"Service": "apigateway.amazonaws.com"},
15+
"Action": "sts:AssumeRole",
16+
}
17+
}
18+
import pytest
19+
20+
from localstack.testing.aws.util import is_aws_cloud
21+
from tests.aws.services.events.helper_functions import is_old_provider
22+
from tests.aws.services.lambda_.test_lambda import (
23+
TEST_LAMBDA_PYTHON_ECHO,
24+
)
25+
26+
27+
@markers.aws.unknown
28+
@pytest.mark.skipif(
29+
condition=is_old_provider() and not is_aws_cloud(),
30+
reason="not supported by the old provider",
31+
)
32+
@markers.snapshot.skip_snapshot_verify(
33+
paths=[
34+
# TODO: those headers are sent by Events via the SDK, we should at least populate X-Amz-Source-Account
35+
# and X-Amz-Source-Arn
36+
"$..headers.amz-sdk-invocation-id",
37+
"$..headers.amz-sdk-request",
38+
"$..headers.amz-sdk-retry",
39+
"$..headers.X-Amz-Security-Token",
40+
"$..headers.X-Amz-Source-Account",
41+
"$..headers.X-Amz-Source-Arn",
42+
# seems like this one can vary in casing between runs?
43+
"$..headers.x-amz-date",
44+
"$..headers.X-Amz-Date",
45+
# those headers are missing in API Gateway
46+
"$..headers.CloudFront-Forwarded-Proto",
47+
"$..headers.CloudFront-Is-Desktop-Viewer",
48+
"$..headers.CloudFront-Is-Mobile-Viewer",
49+
"$..headers.CloudFront-Is-SmartTV-Viewer",
50+
"$..headers.CloudFront-Is-Tablet-Viewer",
51+
"$..headers.CloudFront-Viewer-ASN",
52+
"$..headers.CloudFront-Viewer-Country",
53+
"$..headers.X-Amz-Cf-Id",
54+
"$..headers.Via",
55+
# sent by `requests` library by default
56+
"$..headers.Accept-Encoding",
57+
"$..headers.Accept",
58+
]
59+
)
60+
@markers.snapshot.skip_snapshot_verify(
61+
condition=lambda: not config.APIGW_NEXT_GEN_PROVIDER,
62+
paths=[
63+
# parity issue from previous APIGW implementation
64+
"$..headers.x-localstack-edge",
65+
"$..headers.Connection",
66+
"$..headers.Content-Length",
67+
"$..headers.accept-encoding",
68+
"$..headers.accept",
69+
"$..headers.X-Amzn-Trace-Id",
70+
"$..headers.X-Forwarded-Port",
71+
"$..headers.X-Forwarded-Proto",
72+
"$..pathParameters",
73+
"$..requestContext.authorizer",
74+
"$..requestContext.deploymentId",
75+
"$..requestContext.extendedRequestId",
76+
"$..requestContext.identity",
77+
"$..requestContext.requestId",
78+
"$..stageVariables",
79+
],
80+
)
81+
def test_xray_trace_propagation_events_api_gateway(
82+
aws_client,
83+
create_role_with_policy,
84+
create_lambda_function,
85+
create_rest_apigw,
86+
events_create_event_bus,
87+
events_put_rule,
88+
events_put_targets,
89+
region_name,
90+
account_id,
91+
snapshot,
92+
):
93+
# create lambda
94+
function_name = f"test-function-{short_uid()}"
95+
function_arn = create_lambda_function(
96+
handler_file=TEST_LAMBDA_PYTHON_ECHO,
97+
func_name=function_name,
98+
runtime=Runtime.python3_12,
99+
)["CreateFunctionResponse"]["FunctionArn"]
100+
101+
# create api gateway with lambda integration
102+
# create rest api
103+
api_id, api_name, root = create_rest_apigw(
104+
name=f"test-api-{short_uid()}",
105+
description="Integration test API",
106+
)
107+
108+
resource_id = aws_client.apigateway.create_resource(
109+
restApiId=api_id, parentId=root, pathPart="{proxy+}"
110+
)["id"]
111+
112+
aws_client.apigateway.put_method(
113+
restApiId=api_id,
114+
resourceId=resource_id,
115+
httpMethod="ANY",
116+
authorizationType="NONE",
117+
)
118+
119+
# create role with policy
120+
_, role_arn = create_role_with_policy(
121+
"Allow", "lambda:InvokeFunction", json.dumps(APIGATEWAY_ASSUME_ROLE_POLICY), "*"
122+
)
123+
124+
# Lambda AWS_PROXY integration
125+
aws_client.apigateway.put_integration(
126+
restApiId=api_id,
127+
resourceId=resource_id,
128+
httpMethod="ANY",
129+
type="AWS_PROXY",
130+
integrationHttpMethod="POST",
131+
uri=f"arn:aws:apigateway:us-east-1:lambda:path/2015-03-31/functions/{function_arn}/invocations",
132+
credentials=role_arn,
133+
)
134+
135+
stage_name = "test-api-stage-name"
136+
aws_client.apigateway.create_deployment(restApiId=api_id, stageName=stage_name)
137+
138+
# Create event bus
139+
event_bus_name = f"test-bus-{short_uid()}"
140+
events_create_event_bus(Name=event_bus_name)
141+
142+
# Create rule
143+
rule_name = f"test-rule-{short_uid()}"
144+
event_pattern = {"source": ["test.source"], "detail-type": ["test.detail.type"]}
145+
events_put_rule(
146+
Name=rule_name,
147+
EventBusName=event_bus_name,
148+
EventPattern=json.dumps(event_pattern),
149+
)
150+
151+
# Create an IAM Role for EventBridge to invoke API Gateway
152+
assume_role_policy_document = {
153+
"Version": "2012-10-17",
154+
"Statement": [
155+
{
156+
"Effect": "Allow",
157+
"Principal": {"Service": "events.amazonaws.com"},
158+
"Action": "sts:AssumeRole",
159+
}
160+
],
161+
}
162+
source_arn = f"arn:aws:execute-api:{region_name}:{account_id}:{api_id}/*/POST/test"
163+
role_name, role_arn = create_role_with_policy(
164+
effect="Allow",
165+
actions="execute-api:Invoke",
166+
assume_policy_doc=json.dumps(assume_role_policy_document),
167+
resource=source_arn,
168+
attach=False, # Since we're using put_role_policy, not attach_role_policy
169+
)
170+
171+
# Add the API Gateway as a target with the RoleArn
172+
target_id = f"target-{short_uid()}"
173+
api_target_arn = (
174+
f"arn:aws:execute-api:{region_name}:{account_id}:{api_id}/{stage_name}/POST/test"
175+
)
176+
put_targets_response = aws_client.events.put_targets(
177+
Rule=rule_name,
178+
EventBusName=event_bus_name,
179+
Targets=[
180+
{
181+
"Id": target_id,
182+
"Arn": api_target_arn,
183+
"RoleArn": role_arn,
184+
"Input": json.dumps({"message": "Hello from EventBridge"}),
185+
"RetryPolicy": {"MaximumRetryAttempts": 0},
186+
}
187+
],
188+
)
189+
assert put_targets_response["FailedEntryCount"] == 0
190+
191+
######
192+
# Test
193+
######
194+
event_entry = {
195+
"EventBusName": event_bus_name,
196+
"Source": "test.source",
197+
"DetailType": "test.detail.type",
198+
"Detail": json.dumps({"message": "Hello from EventBridge"}),
199+
}
200+
put_events_response = aws_client.events.put_events(Entries=[event_entry])
201+
snapshot.match("put_events_response", put_events_response)
202+
assert put_events_response["FailedEntryCount"] == 0
203+
204+
# Verify the Lambda invocation
205+
events = retry(
206+
check_expected_lambda_log_events_length,
207+
retries=10,
208+
sleep=10,
209+
sleep_before=10 if is_aws_cloud() else 1,
210+
function_name=function_name,
211+
expected_length=1,
212+
logs_client=aws_client.logs,
213+
)
214+
snapshot.match("lambda_logs", events)
215+
# TODO assert that the X-Ray trace ID is present in the logs
216+
# TODO how to assert X-Ray trace ID correct propagation
217+
218+
219+
# def test_xray_trace_propagation_events_lambda():

0 commit comments

Comments
 (0)