Skip to content
Merged
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
14 changes: 7 additions & 7 deletions src/aws_durable_execution_sdk_python_testing/execution.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
from __future__ import annotations

import json
from dataclasses import replace
from datetime import UTC, datetime
from enum import Enum
Expand Down Expand Up @@ -33,7 +32,6 @@
)
from aws_durable_execution_sdk_python_testing.token import (
CheckpointToken,
CallbackToken,
)


Expand Down Expand Up @@ -96,12 +94,12 @@ def new(input: StartDurableExecutionInput) -> Execution: # noqa: A002
durable_execution_arn=str(uuid4()), start_input=input, operations=[]
)

def to_dict(self) -> dict[str, Any]:
"""Serialize execution to dictionary."""
def to_json_dict(self) -> dict[str, Any]:
"""Serialize execution to JSON-serializable dictionary"""
return {
"DurableExecutionArn": self.durable_execution_arn,
"StartInput": self.start_input.to_dict(),
"Operations": [op.to_dict() for op in self.operations],
"Operations": [op.to_json_dict() for op in self.operations],
"Updates": [update.to_dict() for update in self.updates],
"InvocationCompletions": [
completion.to_dict() for completion in self.invocation_completions
Expand All @@ -115,13 +113,15 @@ def to_dict(self) -> dict[str, Any]:
}

@classmethod
def from_dict(cls, data: dict[str, Any]) -> Execution:
def from_json_dict(cls, data: dict[str, Any]) -> Execution:
"""Deserialize execution from dictionary."""
# Reconstruct start_input
start_input = StartDurableExecutionInput.from_dict(data["StartInput"])

# Reconstruct operations
operations = [Operation.from_dict(op_data) for op_data in data["Operations"]]
operations = [
Operation.from_json_dict(op_data) for op_data in data["Operations"]
]

# Create execution
execution = cls(
Expand Down
4 changes: 1 addition & 3 deletions src/aws_durable_execution_sdk_python_testing/invoker.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,12 +12,10 @@
DurableExecutionInvocationInputWithClient,
DurableExecutionInvocationOutput,
InitialExecutionState,
InvocationStatus,
)

from aws_durable_execution_sdk_python_testing.exceptions import (
DurableFunctionsTestError,
ServiceException,
)
from aws_durable_execution_sdk_python_testing.model import LambdaContext

Expand Down Expand Up @@ -239,7 +237,7 @@ def invoke(
response = client.invoke(
FunctionName=function_name,
InvocationType="RequestResponse", # Synchronous invocation
Payload=json.dumps(input.to_dict(), default=str),
Payload=json.dumps(input.to_json_dict()),
)

# Check HTTP status code
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@

import json
import logging
from datetime import UTC, datetime
from pathlib import Path

from aws_durable_execution_sdk_python_testing.exceptions import (
Expand All @@ -16,30 +15,6 @@
)


class DateTimeEncoder(json.JSONEncoder):
"""Custom JSON encoder that handles datetime objects."""

def default(self, obj):
if isinstance(obj, datetime):
return obj.timestamp()
return super().default(obj)


def datetime_object_hook(obj):
"""JSON object hook to convert unix timestamps back to datetime objects."""
if isinstance(obj, dict):
for key, value in obj.items():
if isinstance(value, int | float) and key.endswith(
("_timestamp", "_time", "Timestamp", "Time")
):
try: # noqa: SIM105
obj[key] = datetime.fromtimestamp(value, tz=UTC)
except (ValueError, OSError):
# Leave as number if not a valid timestamp
pass
return obj


class FileSystemExecutionStore(BaseExecutionStore):
"""File system-based execution store for persistence."""

Expand Down Expand Up @@ -69,10 +44,10 @@ def _get_file_path(self, execution_arn: str) -> Path:
def save(self, execution: Execution) -> None:
"""Save execution to file system."""
file_path = self._get_file_path(execution.durable_execution_arn)
data = execution.to_dict()
data = execution.to_json_dict()

with open(file_path, "w", encoding="utf-8") as f:
json.dump(data, f, indent=2, cls=DateTimeEncoder)
json.dump(data, f, indent=2)

def load(self, execution_arn: str) -> Execution:
"""Load execution from file system."""
Expand All @@ -82,9 +57,9 @@ def load(self, execution_arn: str) -> Execution:
raise ResourceNotFoundException(msg)

with open(file_path, encoding="utf-8") as f:
data = json.load(f, object_hook=datetime_object_hook)
data = json.load(f)

return Execution.from_dict(data)
return Execution.from_json_dict(data)

def update(self, execution: Execution) -> None:
"""Update execution in file system (same as save)."""
Expand All @@ -96,8 +71,8 @@ def list_all(self) -> list[Execution]:
for file_path in self._storage_dir.glob("*.json"):
try:
with open(file_path, encoding="utf-8") as f:
data = json.load(f, object_hook=datetime_object_hook)
executions.append(Execution.from_dict(data))
data = json.load(f)
executions.append(Execution.from_json_dict(data))
except (json.JSONDecodeError, KeyError, OSError) as e:
logging.warning("Skipping corrupted file %s: %s", file_path, e)
continue
Expand Down
16 changes: 3 additions & 13 deletions src/aws_durable_execution_sdk_python_testing/stores/sqlite.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,10 +17,6 @@
from aws_durable_execution_sdk_python_testing.stores.base import (
ExecutionStore,
)
from aws_durable_execution_sdk_python_testing.stores.filesystem import (
DateTimeEncoder,
datetime_object_hook,
)


class SQLiteExecutionStore(ExecutionStore):
Expand Down Expand Up @@ -102,7 +98,7 @@ def save(self, execution: Execution) -> None:
execution_op.end_timestamp.timestamp()
if execution_op.end_timestamp
else None,
json.dumps(execution.to_dict(), cls=DateTimeEncoder),
json.dumps(execution.to_json_dict()),
),
)
except sqlite3.Error as e:
Expand All @@ -125,9 +121,7 @@ def load(self, execution_arn: str) -> Execution:
if not row:
raise ResourceNotFoundException(f"Execution {execution_arn} not found")

return Execution.from_dict(
json.loads(row[0], object_hook=datetime_object_hook)
)
return Execution.from_json_dict(json.loads(row[0]))
except sqlite3.Error as e:
raise RuntimeError(f"Failed to load execution {execution_arn}: {e}") from e
except json.JSONDecodeError as e:
Expand Down Expand Up @@ -222,11 +216,7 @@ def query(
executions: list[Execution] = []
for durable_execution_arn, data in rows:
try:
executions.append(
Execution.from_dict(
json.loads(data, object_hook=datetime_object_hook)
)
)
executions.append(Execution.from_json_dict(json.loads(data)))
except (json.JSONDecodeError, ValueError) as e:
# Log corrupted data but continue with other records
print(
Expand Down
3 changes: 1 addition & 2 deletions tests/execution_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@

from aws_durable_execution_sdk_python_testing.exceptions import (
IllegalStateException,
InvalidParameterValueException,
)
from aws_durable_execution_sdk_python_testing.execution import Execution
from aws_durable_execution_sdk_python_testing.model import StartDurableExecutionInput
Expand Down Expand Up @@ -813,7 +812,7 @@ def test_from_dict_with_none_result():
"aws_durable_execution_sdk_python_testing.model.StartDurableExecutionInput.from_dict"
) as mock_from_dict:
mock_from_dict.return_value = Mock()
execution = Execution.from_dict(data)
execution = Execution.from_json_dict(data)
assert execution.result is None


Expand Down
1 change: 0 additions & 1 deletion tests/how-to-run-from-term.txt

This file was deleted.

26 changes: 24 additions & 2 deletions tests/invoker_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,15 @@
InvocationStatus,
)

from aws_durable_execution_sdk_python.lambda_service import (
ExecutionDetails,
Operation,
OperationStatus,
OperationType,
)

from datetime import datetime, UTC

from aws_durable_execution_sdk_python_testing.execution import Execution
from aws_durable_execution_sdk_python_testing.invoker import (
InProcessInvoker,
Expand Down Expand Up @@ -168,10 +177,23 @@ def test_lambda_invoker_invoke_success():

invoker = LambdaInvoker(lambda_client)

mock_operation = Operation(
operation_id="op-1",
parent_id=None,
name="test-execution",
start_timestamp=datetime.now(UTC),
end_timestamp=datetime.now(UTC),
operation_type=OperationType.EXECUTION,
status=OperationStatus.SUCCEEDED,
execution_details=ExecutionDetails(input_payload='{"test": "data"}'),
)

input_data = DurableExecutionInvocationInput(
durable_execution_arn="test-arn",
checkpoint_token="test-token", # noqa: S106
initial_execution_state=InitialExecutionState(operations=[], next_marker=""),
initial_execution_state=InitialExecutionState(
operations=[mock_operation], next_marker=""
),
)

response = invoker.invoke("test-function", input_data)
Expand All @@ -185,7 +207,7 @@ def test_lambda_invoker_invoke_success():
lambda_client.invoke.assert_called_once_with(
FunctionName="test-function",
InvocationType="RequestResponse",
Payload=json.dumps(input_data.to_dict(), default=str),
Payload=json.dumps(input_data.to_json_dict()),
)


Expand Down
14 changes: 0 additions & 14 deletions tests/stores/filesystem_store_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,6 @@
from aws_durable_execution_sdk_python_testing.model import StartDurableExecutionInput
from aws_durable_execution_sdk_python_testing.stores.filesystem import (
FileSystemExecutionStore,
datetime_object_hook,
)

from datetime import datetime, timezone
Expand Down Expand Up @@ -269,19 +268,6 @@ def test_filesystem_execution_store_thread_safety_basic(store, sample_execution)
assert loaded.durable_execution_arn == sample_execution.durable_execution_arn


def test_datetime_object_hook_converts_timestamp_fields():
"""Test conversion of timestamp fields to datetime objects."""
timestamp = 1672531200.0 # 2023-01-01 00:00:00 UTC
obj = {
"start_timestamp": timestamp,
}

result = datetime_object_hook(obj)

expected_datetime = datetime.fromtimestamp(timestamp, tz=timezone.utc)
assert result["start_timestamp"] == expected_datetime


def test_filesystem_execution_store_query_empty(store):
"""Test query method with empty store."""
executions, next_marker = store.query()
Expand Down
Loading