refactor: use to/from_json_dict for operation serdes

- Remove customized json encoder
- Use to_json_dict() and from_json_dict() for Operation and ExecutionInput serdes

Author: Alex Wang

src/aws_durable_execution_sdk_python_testing/execution.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import json
 from dataclasses import replace
 from datetime import UTC, datetime
 from enum import Enum
@@ -33,7 +32,6 @@
 )
 from aws_durable_execution_sdk_python_testing.token import (
     CheckpointToken,
-    CallbackToken,
 )
 
 
@@ -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
@@ -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(

src/aws_durable_execution_sdk_python_testing/invoker.py

@@ -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
 
@@ -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

src/aws_durable_execution_sdk_python_testing/stores/filesystem.py

@@ -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 (
@@ -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."""
 
@@ -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."""
@@ -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)."""
@@ -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

src/aws_durable_execution_sdk_python_testing/stores/sqlite.py

@@ -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):
@@ -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:
@@ -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:
@@ -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(

tests/execution_test.py

@@ -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
@@ -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
 
 

tests/how-to-run-from-term.txt

@@ -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,
@@ -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)
@@ -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()),
     )
 
 

tests/invoker_test.py

@@ -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
@@ -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()