Skip to content

Commit 6df9b1d

Browse files
heitorlessaheitorlessa
authored
refactor(e2e): support fail fast in get_lambda_response (aws-powertools#2912)
1 parent 6177b93 commit 6df9b1d

4 files changed

Lines changed: 110 additions & 44 deletions

File tree

‎tests/e2e/idempotency/test_idempotency_dynamodb.py‎

Lines changed: 18 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
import pytest
55

66
from tests.e2e.utils import data_fetcher
7-
from tests.e2e.utils.functions import execute_lambdas_in_parallel
7+
from tests.e2e.utils.data_fetcher.common import GetLambdaResponseOptions, get_lambda_response_in_parallel
88

99

1010
@pytest.fixture
@@ -73,14 +73,21 @@ def test_ttl_caching_expiration_idempotency(ttl_cache_expiration_handler_fn_arn:
7373
@pytest.mark.xdist_group(name="idempotency")
7474
def test_ttl_caching_timeout_idempotency(ttl_cache_timeout_handler_fn_arn: str):
7575
# GIVEN
76-
payload_timeout_execution = json.dumps({"sleep": 5, "message": "Powertools for AWS Lambda (Python) - TTL 1s"})
77-
payload_working_execution = json.dumps({"sleep": 0, "message": "Powertools for AWS Lambda (Python) - TTL 1s"})
76+
payload_timeout_execution = json.dumps(
77+
{"sleep": 5, "message": "Powertools for AWS Lambda (Python) - TTL 1s"},
78+
sort_keys=True,
79+
)
80+
payload_working_execution = json.dumps(
81+
{"sleep": 0, "message": "Powertools for AWS Lambda (Python) - TTL 1s"},
82+
sort_keys=True,
83+
)
7884

7985
# WHEN
8086
# first call should fail due to timeout
8187
execution_with_timeout, _ = data_fetcher.get_lambda_response(
8288
lambda_arn=ttl_cache_timeout_handler_fn_arn,
8389
payload=payload_timeout_execution,
90+
raise_on_error=False,
8491
)
8592
execution_with_timeout_response = execution_with_timeout["Payload"].read().decode("utf-8")
8693

@@ -99,12 +106,15 @@ def test_ttl_caching_timeout_idempotency(ttl_cache_timeout_handler_fn_arn: str):
99106
@pytest.mark.xdist_group(name="idempotency")
100107
def test_parallel_execution_idempotency(parallel_execution_handler_fn_arn: str):
101108
# GIVEN
102-
arguments = json.dumps({"message": "Powertools for AWS Lambda (Python) - Parallel execution"})
109+
payload = json.dumps({"message": "Powertools for AWS Lambda (Python) - Parallel execution"})
103110

104-
# WHEN
105-
# executing Lambdas in parallel
106-
lambdas_arn = [parallel_execution_handler_fn_arn, parallel_execution_handler_fn_arn]
107-
execution_result_list = execute_lambdas_in_parallel("data_fetcher.get_lambda_response", lambdas_arn, arguments)
111+
invocation_options = [
112+
GetLambdaResponseOptions(lambda_arn=parallel_execution_handler_fn_arn, payload=payload, raise_on_error=False),
113+
GetLambdaResponseOptions(lambda_arn=parallel_execution_handler_fn_arn, payload=payload, raise_on_error=False),
114+
]
115+
116+
# WHEN executing Lambdas in parallel
117+
execution_result_list = get_lambda_response_in_parallel(invocation_options)
108118

109119
timeout_execution_response = execution_result_list[0][0]["Payload"].read().decode("utf-8")
110120
error_idempotency_execution_response = execution_result_list[1][0]["Payload"].read().decode("utf-8")

‎tests/e2e/utils/data_builder/traces.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33

44
def build_trace_default_query(function_name: str) -> str:
5-
return f'service("{function_name}")'
5+
return f'service(id(name: "{function_name}"))'
66

77

88
def build_put_annotations_input(**annotations: str) -> List[Dict]:

‎tests/e2e/utils/data_fetcher/common.py‎

Lines changed: 91 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,24 +1,76 @@
1+
import functools
2+
import time
3+
from concurrent.futures import Future, ThreadPoolExecutor
14
from datetime import datetime
2-
from typing import Optional, Tuple
5+
from typing import List, Optional, Tuple
36

47
import boto3
58
import requests
69
from mypy_boto3_lambda import LambdaClient
710
from mypy_boto3_lambda.type_defs import InvocationResponseTypeDef
11+
from pydantic import BaseModel
812
from requests import Request, Response
913
from requests.exceptions import RequestException
1014
from retry import retry
1115

16+
GetLambdaResponse = Tuple[InvocationResponseTypeDef, datetime]
17+
18+
19+
class GetLambdaResponseOptions(BaseModel):
20+
lambda_arn: str
21+
payload: Optional[str] = None
22+
client: Optional[LambdaClient] = None
23+
raise_on_error: bool = True
24+
25+
# Maintenance: Pydantic v2 deprecated it; we should update in v3
26+
class Config:
27+
arbitrary_types_allowed = True
28+
1229

1330
def get_lambda_response(
1431
lambda_arn: str,
1532
payload: Optional[str] = None,
1633
client: Optional[LambdaClient] = None,
17-
) -> Tuple[InvocationResponseTypeDef, datetime]:
34+
raise_on_error: bool = True,
35+
) -> GetLambdaResponse:
36+
"""Invoke function synchronously
37+
38+
Parameters
39+
----------
40+
lambda_arn : str
41+
Lambda function ARN to invoke
42+
payload : Optional[str], optional
43+
JSON payload for Lambda invocation, by default None
44+
client : Optional[LambdaClient], optional
45+
Boto3 Lambda SDK client, by default None
46+
raise_on_error : bool, optional
47+
Whether to raise exception upon invocation error, by default True
48+
49+
Returns
50+
-------
51+
Tuple[InvocationResponseTypeDef, datetime]
52+
Function response and approximate execution time
53+
54+
Raises
55+
------
56+
RuntimeError
57+
Function invocation error details
58+
"""
1859
client = client or boto3.client("lambda")
1960
payload = payload or ""
2061
execution_time = datetime.utcnow()
21-
return client.invoke(FunctionName=lambda_arn, InvocationType="RequestResponse", Payload=payload), execution_time
62+
response: InvocationResponseTypeDef = client.invoke(
63+
FunctionName=lambda_arn,
64+
InvocationType="RequestResponse",
65+
Payload=payload,
66+
)
67+
68+
has_error = response.get("FunctionError", "") == "Unhandled"
69+
if has_error and raise_on_error:
70+
error_payload = response["Payload"].read().decode()
71+
raise RuntimeError(f"Function failed invocation: {error_payload}")
72+
73+
return response, execution_time
2274

2375

2476
@retry(RequestException, delay=2, jitter=1.5, tries=5)
@@ -27,3 +79,39 @@ def get_http_response(request: Request) -> Response:
2779
result = session.send(request.prepare())
2880
result.raise_for_status()
2981
return result
82+
83+
84+
def get_lambda_response_in_parallel(
85+
get_lambda_response_options: List[GetLambdaResponseOptions],
86+
) -> List[GetLambdaResponse]:
87+
"""Invoke functions in parallel
88+
89+
Parameters
90+
----------
91+
get_lambda_response_options : List[GetLambdaResponseOptions]
92+
List of options to call get_lambda_response with
93+
94+
Returns
95+
-------
96+
List[GetLambdaResponse]
97+
Function responses and approximate execution time
98+
"""
99+
result_list = []
100+
with ThreadPoolExecutor() as executor:
101+
running_tasks: List[Future] = []
102+
for options in get_lambda_response_options:
103+
# Sleep 0.5, 1, 1.5, ... seconds between each invocation. This way
104+
# we can guarantee that lambdas are executed in parallel, but they are
105+
# called in the same "order" as they are passed in, thus guaranteeing that
106+
# we can assert on the correct output.
107+
time.sleep(0.5 * len(running_tasks))
108+
109+
get_lambda_response_callback = functools.partial(get_lambda_response, **options.dict())
110+
running_tasks.append(
111+
executor.submit(get_lambda_response_callback),
112+
)
113+
114+
executor.shutdown(wait=True)
115+
result_list.extend(running_task.result() for running_task in running_tasks)
116+
117+
return result_list

‎tests/e2e/utils/functions.py‎

Lines changed: 0 additions & 32 deletions
This file was deleted.

0 commit comments

Comments
 (0)