Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 27 additions & 9 deletions src/aws_durable_execution_sdk_python_testing/invoker.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,8 @@
from uuid import uuid4

import boto3 # type: ignore
from botocore.config import Config # type: ignore

from aws_durable_execution_sdk_python.execution import (
DurableExecutionInvocationInput,
DurableExecutionInvocationInputWithClient,
Expand All @@ -29,6 +31,26 @@
from aws_durable_execution_sdk_python_testing.execution import Execution


# Max Lambda function timeout is 15 minutes (900s); we give headroom for
# network round-trip and RIE startup.
_LAMBDA_READ_TIMEOUT_SECONDS = 960
_LAMBDA_CLIENT_CONFIG = Config(
read_timeout=_LAMBDA_READ_TIMEOUT_SECONDS,
retries={"max_attempts": 0},
)


def create_lambda_client(endpoint_url: str | None, region_name: str) -> Any:
"""Create a boto3 Lambda client configured for durable function invocations."""

return boto3.client(
"lambda",
endpoint_url=endpoint_url,
region_name=region_name,
config=_LAMBDA_CLIENT_CONFIG,
)


@dataclass(frozen=True)
class InvokeResponse:
"""Response from invoking a durable function."""
Expand Down Expand Up @@ -136,9 +158,7 @@ def __init__(self, lambda_client: Any) -> None:
@staticmethod
def create(endpoint_url: str, region_name: str) -> LambdaInvoker:
"""Create with the boto lambda client."""
invoker = LambdaInvoker(
boto3.client("lambda", endpoint_url=endpoint_url, region_name=region_name)
)
invoker = LambdaInvoker(create_lambda_client(endpoint_url, region_name))
invoker._current_endpoint = endpoint_url
invoker._endpoint_clients[endpoint_url] = invoker.lambda_client
return invoker
Expand All @@ -148,8 +168,8 @@ def update_endpoint(self, endpoint_url: str, region_name: str) -> None:
# Cache client by endpoint to reuse across executions
with self._lock:
if endpoint_url not in self._endpoint_clients:
self._endpoint_clients[endpoint_url] = boto3.client(
"lambda", endpoint_url=endpoint_url, region_name=region_name
self._endpoint_clients[endpoint_url] = create_lambda_client(
endpoint_url, region_name
)
self.lambda_client = self._endpoint_clients[endpoint_url]
self._current_endpoint = endpoint_url
Expand All @@ -164,10 +184,8 @@ def _get_client_for_execution(
# Use provided endpoint or fall back to cached endpoint for this execution
if lambda_endpoint:
if lambda_endpoint not in self._endpoint_clients:
self._endpoint_clients[lambda_endpoint] = boto3.client(
"lambda",
endpoint_url=lambda_endpoint,
region_name=region_name or "us-east-1",
self._endpoint_clients[lambda_endpoint] = create_lambda_client(
lambda_endpoint, region_name or "us-east-1"
)
return self._endpoint_clients[lambda_endpoint]

Expand Down
5 changes: 2 additions & 3 deletions src/aws_durable_execution_sdk_python_testing/runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@
from aws_durable_execution_sdk_python_testing.invoker import (
InProcessInvoker,
LambdaInvoker,
create_lambda_client,
)
from aws_durable_execution_sdk_python_testing.model import (
GetDurableExecutionHistoryResponse,
Expand Down Expand Up @@ -858,9 +859,7 @@ def _create_boto3_client(self) -> Any:
Exception: If client creation fails - exceptions propagate naturally
for CLI to handle as general Exception
"""
# Create client with Lambda endpoint configuration
return boto3.client(
"lambda",
return create_lambda_client(
endpoint_url=self._config.lambda_endpoint,
region_name=self._config.local_runner_region,
)
Expand Down
3 changes: 3 additions & 0 deletions tests/invoker_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,8 @@
from aws_durable_execution_sdk_python_testing.invoker import (
InProcessInvoker,
LambdaInvoker,
_LAMBDA_CLIENT_CONFIG,
create_lambda_client,
create_test_lambda_context,
)
from aws_durable_execution_sdk_python_testing.model import (
Expand Down Expand Up @@ -133,6 +135,7 @@ def test_lambda_invoker_create():
"lambda",
endpoint_url="http://localhost:3001",
region_name="us-west-2",
config=_LAMBDA_CLIENT_CONFIG,
)


Expand Down
4 changes: 4 additions & 0 deletions tests/runner_web_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
from aws_durable_execution_sdk_python_testing.exceptions import (
DurableFunctionsLocalRunnerError,
)
from aws_durable_execution_sdk_python_testing.invoker import _LAMBDA_CLIENT_CONFIG
from aws_durable_execution_sdk_python_testing.runner import (
WebRunner,
WebRunnerConfig,
Expand Down Expand Up @@ -420,6 +421,7 @@ def test_should_handle_boto3_client_creation_with_custom_config():
"lambda",
endpoint_url="http://custom-endpoint:8080",
region_name="eu-west-1",
config=_LAMBDA_CLIENT_CONFIG,
)

# Verify public behavior works
Expand All @@ -445,6 +447,7 @@ def test_should_handle_boto3_client_creation_with_defaults():
"lambda",
endpoint_url="http://127.0.0.1:3001", # Default lambda_endpoint value
region_name="us-west-2", # Default value
config=_LAMBDA_CLIENT_CONFIG,
)

# Verify public behavior works
Expand Down Expand Up @@ -768,6 +771,7 @@ def test_should_pass_correct_boto3_client_to_lambda_invoker():
"lambda",
endpoint_url="http://test-endpoint:7777",
region_name="ap-southeast-2",
config=_LAMBDA_CLIENT_CONFIG,
)

# Verify LambdaInvoker was created with the client
Expand Down
Loading