Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
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
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ classifiers = [
"Programming Language :: Python :: Implementation :: CPython",
"Programming Language :: Python :: Implementation :: PyPy",
]
dependencies = ["boto3>=1.40.30"]
dependencies = ["boto3>=1.42.1"]

[project.urls]
Documentation = "https://github.com/aws/aws-durable-execution-sdk-python#readme"
Expand Down

This file was deleted.

2 changes: 1 addition & 1 deletion src/aws_durable_execution_sdk_python/execution.py
Original file line number Diff line number Diff line change
Expand Up @@ -237,7 +237,7 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]:
service_client = (
LambdaClient(client=boto3_client)
if boto3_client is not None
else LambdaClient.initialize_from_env()
else LambdaClient.initialize_client()
)

raw_input_payload: str | None = (
Expand Down
43 changes: 7 additions & 36 deletions src/aws_durable_execution_sdk_python/lambda_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,8 @@

import datetime
import logging
import os
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import TYPE_CHECKING, Any, Protocol, TypeAlias

import boto3 # type: ignore
Expand Down Expand Up @@ -943,41 +941,14 @@ def __init__(self, client: Any) -> None:
self.client = client

@staticmethod
def load_preview_botocore_models() -> None:
"""
Load boto3 models from the Python path for custom preview client.
"""
os.environ["AWS_DATA_PATH"] = str(
Path(__file__).parent.joinpath("botocore", "data")
def initialize_client() -> LambdaClient:
client = boto3.client(
"lambda",
config=Config(
connect_timeout=5,
read_timeout=50,
),
)

@staticmethod
def initialize_from_env() -> LambdaClient:
LambdaClient.load_preview_botocore_models()

"""
TODO - we can remove this when were using the actual lambda client,
but we need this with the preview model because boto won't match against lambdainternal.
"""
endpoint_url = os.getenv("AWS_ENDPOINT_URL_LAMBDA", None)
if not endpoint_url:
client = boto3.client(
"lambdainternal",
config=Config(
connect_timeout=5,
read_timeout=50,
),
)
else:
client = boto3.client(
"lambdainternal",
endpoint_url=endpoint_url,
config=Config(
connect_timeout=5,
read_timeout=50,
),
)

return LambdaClient(client=client)

def checkpoint(
Expand Down
12 changes: 6 additions & 6 deletions tests/e2e/execution_int_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,7 @@ def my_handler(event, context: DurableContext) -> list[str]:
"aws_durable_execution_sdk_python.execution.LambdaClient"
) as mock_client_class:
mock_client = Mock()
mock_client_class.initialize_from_env.return_value = mock_client
mock_client_class.initialize_client.return_value = mock_client

# Mock the checkpoint method to track calls
checkpoint_calls = []
Expand Down Expand Up @@ -156,7 +156,7 @@ def my_handler(event, context: DurableContext):
"aws_durable_execution_sdk_python.execution.LambdaClient"
) as mock_client_class:
mock_client = Mock()
mock_client_class.initialize_from_env.return_value = mock_client
mock_client_class.initialize_client.return_value = mock_client

# Mock the checkpoint method to track calls
checkpoint_calls = []
Expand Down Expand Up @@ -257,7 +257,7 @@ def my_handler(event, context):
"aws_durable_execution_sdk_python.execution.LambdaClient"
) as mock_client_class:
mock_client = Mock()
mock_client_class.initialize_from_env.return_value = mock_client
mock_client_class.initialize_client.return_value = mock_client

# Mock the checkpoint method to track calls
checkpoint_calls = []
Expand Down Expand Up @@ -363,7 +363,7 @@ def my_handler(event, context: DurableContext):
"aws_durable_execution_sdk_python.execution.LambdaClient"
) as mock_client_class:
mock_client = Mock()
mock_client_class.initialize_from_env.return_value = mock_client
mock_client_class.initialize_client.return_value = mock_client

# Mock the checkpoint method to raise an error (using RuntimeError as a generic exception)
def mock_checkpoint_failure(
Expand Down Expand Up @@ -426,7 +426,7 @@ def my_handler(event: Any, context: DurableContext):
"aws_durable_execution_sdk_python.execution.LambdaClient"
) as mock_client_class:
mock_client = Mock()
mock_client_class.initialize_from_env.return_value = mock_client
mock_client_class.initialize_client.return_value = mock_client

# Mock the checkpoint method to track calls
checkpoint_calls = []
Expand Down Expand Up @@ -509,7 +509,7 @@ def my_handler(event, context):
"aws_durable_execution_sdk_python.execution.LambdaClient"
) as mock_client_class:
mock_client = Mock()
mock_client_class.initialize_from_env.return_value = mock_client
mock_client_class.initialize_client.return_value = mock_client

checkpoint_calls = []

Expand Down
12 changes: 6 additions & 6 deletions tests/execution_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -337,7 +337,7 @@ def test_durable_execution_client_selection_env_normal_result():
"aws_durable_execution_sdk_python.execution.LambdaClient"
) as mock_lambda_client:
mock_client = Mock(spec=DurableServiceClient)
mock_lambda_client.initialize_from_env.return_value = mock_client
mock_lambda_client.initialize_client.return_value = mock_client

# Mock successful checkpoint
mock_output = CheckpointOutput(
Expand Down Expand Up @@ -379,7 +379,7 @@ def test_handler(event: Any, context: DurableContext) -> dict:

assert result["Status"] == InvocationStatus.SUCCEEDED.value
assert result["Result"] == '{"result": "success"}'
mock_lambda_client.initialize_from_env.assert_called_once()
mock_lambda_client.initialize_client.assert_called_once()
mock_client.checkpoint.assert_not_called()


Expand All @@ -389,7 +389,7 @@ def test_durable_execution_client_selection_env_large_result():
"aws_durable_execution_sdk_python.execution.LambdaClient"
) as mock_lambda_client:
mock_client = Mock(spec=DurableServiceClient)
mock_lambda_client.initialize_from_env.return_value = mock_client
mock_lambda_client.initialize_client.return_value = mock_client

# Mock successful checkpoint
mock_output = CheckpointOutput(
Expand Down Expand Up @@ -431,7 +431,7 @@ def test_handler(event: Any, context: DurableContext) -> dict:

assert result["Status"] == InvocationStatus.SUCCEEDED.value
assert not result["Result"]
mock_lambda_client.initialize_from_env.assert_called_once()
mock_lambda_client.initialize_client.assert_called_once()
mock_client.checkpoint.assert_called_once()


Expand Down Expand Up @@ -725,7 +725,7 @@ def test_durable_execution_client_selection_default():
"aws_durable_execution_sdk_python.execution.LambdaClient"
) as mock_lambda_client:
mock_client = Mock(spec=DurableServiceClient)
mock_lambda_client.initialize_from_env.return_value = mock_client
mock_lambda_client.initialize_client.return_value = mock_client

# Mock successful checkpoint
mock_output = CheckpointOutput(
Expand Down Expand Up @@ -766,7 +766,7 @@ def test_handler(event: Any, context: DurableContext) -> dict:
result = test_handler(event, lambda_context)

assert result["Status"] == InvocationStatus.SUCCEEDED.value
mock_lambda_client.initialize_from_env.assert_called_once()
mock_lambda_client.initialize_client.assert_called_once()


def test_initial_execution_state_get_execution_operation_no_operations():
Expand Down
47 changes: 17 additions & 30 deletions tests/lambda_service_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -1906,18 +1906,17 @@ def test_lambda_client_constructor():

@patch.dict("os.environ", {}, clear=True)
@patch("boto3.client")
def test_lambda_client_initialize_from_env_default(mock_boto_client):
"""Test LambdaClient.initialize_from_env with default endpoint."""
def test_lambda_client_initialize_client_default(mock_boto_client):
"""Test LambdaClient.initialize_client with default endpoint."""
mock_client = Mock()
mock_boto_client.return_value = mock_client

with patch.object(LambdaClient, "load_preview_botocore_models"):
client = LambdaClient.initialize_from_env()
client = LambdaClient.initialize_client()

# Check that boto3.client was called with the right service name and config
mock_boto_client.assert_called_once()
call_args = mock_boto_client.call_args
assert call_args[0][0] == "lambdainternal"
assert call_args[0][0] == "lambda"
assert "config" in call_args[1]
config = call_args[1]["config"]
assert config.connect_timeout == 5
Expand All @@ -1927,19 +1926,18 @@ def test_lambda_client_initialize_from_env_default(mock_boto_client):

@patch.dict("os.environ", {"AWS_ENDPOINT_URL_LAMBDA": "http://localhost:3000"})
@patch("boto3.client")
def test_lambda_client_initialize_from_env_with_endpoint(mock_boto_client):
"""Test LambdaClient.initialize_from_env with custom endpoint."""
def test_lambda_client_initialize_client_with_endpoint(mock_boto_client):
"""Test LambdaClient.initialize_client with custom endpoint (boto3 handles it automatically)."""
mock_client = Mock()
mock_boto_client.return_value = mock_client

with patch.object(LambdaClient, "load_preview_botocore_models"):
client = LambdaClient.initialize_from_env()
client = LambdaClient.initialize_client()

# Check that boto3.client was called with the right parameters and config
# Note: boto3 automatically picks up AWS_ENDPOINT_URL_LAMBDA from environment
mock_boto_client.assert_called_once()
call_args = mock_boto_client.call_args
assert call_args[0][0] == "lambdainternal"
assert call_args[1]["endpoint_url"] == "http://localhost:3000"
assert call_args[0][0] == "lambda"
assert "config" in call_args[1]
config = call_args[1]["config"]
assert config.connect_timeout == 5
Expand Down Expand Up @@ -1981,23 +1979,13 @@ def test_durable_service_client_protocol_get_execution_state():


@patch.dict("os.environ", {}, clear=True)
@patch(
"aws_durable_execution_sdk_python.lambda_service.LambdaClient.initialize_from_env"
)
def test_lambda_client_initialize_from_env_defaults(mock_init):
"""Test LambdaClient.initialize_from_env with default environment values."""
LambdaClient.initialize_from_env()
@patch("aws_durable_execution_sdk_python.lambda_service.LambdaClient.initialize_client")
def test_lambda_client_initialize_client_defaults(mock_init):
"""Test LambdaClient.initialize_client with default environment values."""
LambdaClient.initialize_client()
mock_init.assert_called_once_with()


@patch("os.environ")
def test_lambda_client_load_preview_botocore_models(mock_environ):
"""Test LambdaClient.load_preview_botocore_models method."""
LambdaClient.load_preview_botocore_models()
# Verify that AWS_DATA_PATH is set
assert "AWS_DATA_PATH" in mock_environ.__setitem__.call_args[0]


def test_checkpoint_error_handling():
"""Test CheckpointError exception handling in LambdaClient.checkpoint."""
mock_client = Mock()
Expand All @@ -2016,17 +2004,16 @@ def test_checkpoint_error_handling():

@patch.dict("os.environ", {}, clear=True)
@patch("boto3.client")
def test_lambda_client_initialize_from_env_no_endpoint(mock_boto_client):
"""Test LambdaClient.initialize_from_env without AWS_ENDPOINT_URL_LAMBDA."""
def test_lambda_client_initialize_client_no_endpoint(mock_boto_client):
"""Test LambdaClient.initialize_client without AWS_ENDPOINT_URL_LAMBDA."""
mock_client = Mock()
mock_boto_client.return_value = mock_client

with patch.object(LambdaClient, "load_preview_botocore_models"):
client = LambdaClient.initialize_from_env()
client = LambdaClient.initialize_client()

# Verify the call was made with the expected arguments including config
call_args = mock_boto_client.call_args
assert call_args[0] == ("lambdainternal",)
assert call_args[0] == ("lambda",)
assert "config" in call_args[1]
assert isinstance(client, LambdaClient)

Expand Down
Loading