fix: preserve step retry attempt

• Fix checkpoint processor to preserve attempt count and timestamp
• Add DurableChildContextTestRunner

Author: yaythomas

src/aws_durable_execution_sdk_python_testing/checkpoint/processors/base.py

@@ -70,13 +70,27 @@ def _create_context_details(self, update: OperationUpdate) -> ContextDetails | N
             else None
         )
 
-    def _create_step_details(self, update: OperationUpdate) -> StepDetails | None:
+    def _create_step_details(
+        self, update: OperationUpdate, current_operation: Operation | None = None
+    ) -> StepDetails | None:
         """Create StepDetails from OperationUpdate."""
-        return (
-            StepDetails(result=update.payload, error=update.error)
-            if update.operation_type == OperationType.STEP
-            else None
-        )
+        attempt: int = 0
+        next_attempt_timestamp: str | None = None
+
+        if update.operation_type is OperationType.STEP:
+            if current_operation and current_operation.step_details:
+                attempt = current_operation.step_details.attempt
+                next_attempt_timestamp = (
+                    current_operation.step_details.next_attempt_timestamp
+                )
+            return StepDetails(
+                attempt=attempt,
+                next_attempt_timestamp=next_attempt_timestamp,
+                result=update.payload,
+                error=update.error,
+            )
+
+        return None
 
     def _create_callback_details(
         self, update: OperationUpdate
@@ -134,7 +148,7 @@ def _translate_update_to_operation(
 
         execution_details = self._create_execution_details(update)
         context_details = self._create_context_details(update)
-        step_details = self._create_step_details(update)
+        step_details = self._create_step_details(update, current_operation)
         callback_details = self._create_callback_details(update)
         invoke_details = self._create_invoke_details(update)
         wait_details = self._create_wait_details(update, current_operation)

src/aws_durable_execution_sdk_python_testing/observer.py

@@ -1,10 +1,15 @@
 """Checkpoint processors can notify the Execution of notable event state changes. Observer pattern."""
 
+from __future__ import annotations
+
 import threading
 from abc import ABC, abstractmethod
-from collections.abc import Callable
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
 
-from aws_durable_execution_sdk_python.lambda_service import ErrorObject
+    from aws_durable_execution_sdk_python.lambda_service import ErrorObject
 
 
 class ExecutionObserver(ABC):
@@ -34,7 +39,7 @@ def on_step_retry_scheduled(
 class ExecutionNotifier:
     """Notifies observers about execution events. Thread-safe."""
 
-    def __init__(self):
+    def __init__(self) -> None:
         self._observers: list[ExecutionObserver] = []
         self._lock = threading.RLock()
 

src/aws_durable_execution_sdk_python_testing/runner.py

@@ -1,8 +1,13 @@
 from __future__ import annotations
 
+import json
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Protocol, TypeVar, cast
+from typing import TYPE_CHECKING, Any, Concatenate, ParamSpec, Protocol, TypeVar, cast
 
+from aws_durable_execution_sdk_python.execution import (
+    InvocationStatus,
+    durable_handler,
+)
 from aws_durable_execution_sdk_python.lambda_service import (
     ErrorObject,
     OperationStatus,
@@ -31,6 +36,7 @@
     import datetime
     from collections.abc import Callable, MutableMapping
 
+    from aws_durable_execution_sdk_python.context import DurableContext
     from aws_durable_execution_sdk_python.execution import InvocationStatus
 
     from aws_durable_execution_sdk_python_testing.execution import Execution
@@ -49,6 +55,7 @@ class Operation:
 
 
 T = TypeVar("T", bound=Operation)
+P = ParamSpec("P")
 
 
 class OperationFactory(Protocol):
@@ -90,7 +97,7 @@ def from_svc_operation(
 @dataclass(frozen=True)
 class ContextOperation(Operation):
     child_operations: list[Operation]
-    result: str | None = None
+    result: Any = None
     error: ErrorObject | None = None
 
     @staticmethod
@@ -119,9 +126,11 @@ def from_svc_operation(
             start_timestamp=operation.start_timestamp,
             end_timestamp=operation.end_timestamp,
             child_operations=child_operations,
-            result=operation.context_details.result
-            if operation.context_details
-            else None,
+            result=(
+                json.loads(operation.context_details.result)
+                if operation.context_details and operation.context_details.result
+                else None
+            ),
             error=operation.context_details.error
             if operation.context_details
             else None,
@@ -157,8 +166,7 @@ def get_execution(self, name: str) -> ExecutionOperation:
 class StepOperation(ContextOperation):
     attempt: int = 0
     next_attempt_timestamp: str | None = None
-    # TODO: deserialize?
-    result: str | None = None
+    result: Any = None
     error: ErrorObject | None = None
 
     @staticmethod
@@ -193,7 +201,11 @@ def from_svc_operation(
                 if operation.step_details
                 else None
             ),
-            result=operation.step_details.result if operation.step_details else None,
+            result=(
+                json.loads(operation.step_details.result)
+                if operation.step_details and operation.step_details.result
+                else None
+            ),
             error=operation.step_details.error if operation.step_details else None,
         )
 
@@ -230,7 +242,7 @@ def from_svc_operation(
 @dataclass(frozen=True)
 class CallbackOperation(ContextOperation):
     callback_id: str | None = None
-    result: str | None = None
+    result: Any = None
     error: ErrorObject | None = None
 
     @staticmethod
@@ -264,9 +276,11 @@ def from_svc_operation(
                 if operation.callback_details
                 else None
             ),
-            result=operation.callback_details.result
-            if operation.callback_details
-            else None,
+            result=(
+                json.loads(operation.callback_details.result)
+                if operation.callback_details and operation.callback_details.result
+                else None
+            ),
             error=operation.callback_details.error
             if operation.callback_details
             else None,
@@ -276,7 +290,7 @@ def from_svc_operation(
 @dataclass(frozen=True)
 class InvokeOperation(Operation):
     durable_execution_arn: str | None = None
-    result: str | None = None
+    result: Any = None
     error: ErrorObject | None = None
 
     @staticmethod
@@ -301,9 +315,11 @@ def from_svc_operation(
                 if operation.invoke_details
                 else None
             ),
-            result=operation.invoke_details.result
-            if operation.invoke_details
-            else None,
+            result=(
+                json.loads(operation.invoke_details.result)
+                if operation.invoke_details and operation.invoke_details.result
+                else None
+            ),
             error=operation.invoke_details.error if operation.invoke_details else None,
         )
 
@@ -334,7 +350,7 @@ def create_operation(
 class DurableFunctionTestResult:
     status: InvocationStatus
     operations: list[Operation]
-    result: str | None = None
+    result: Any = None
     error: ErrorObject | None = None
 
     @classmethod
@@ -352,10 +368,14 @@ def create(cls, execution: Execution) -> DurableFunctionTestResult:
             msg: str = "Execution result must exist to create test result."
             raise DurableFunctionsTestError(msg)
 
+        deserialized_result = (
+            json.loads(execution.result.result) if execution.result.result else None
+        )
+
         return cls(
             status=execution.result.status,
             operations=operations,
-            result=execution.result.result,
+            result=deserialized_result,
             error=execution.result.error,
         )
 
@@ -413,7 +433,7 @@ def close(self):
 
     def run(
         self,
-        input: str,  # noqa: A002
+        input: str | None = None,  # noqa: A002
         timeout: int = 900,
         function_name: str = "test-function",
         execution_name: str = "execution-name",
@@ -451,4 +471,19 @@ def run(
         execution: Execution = self._store.load(output.execution_arn)
         return DurableFunctionTestResult.create(execution=execution)
 
-        # return execution
+
+class DurableChildContextTestRunner(DurableFunctionTestRunner):
+    """Test a durable block, annotated with @durable_with_child_context, in isolation."""
+
+    def __init__(
+        self,
+        context_function: Callable[Concatenate[DurableContext, P], Any],
+        *args,
+        **kwargs,
+    ):
+        # wrap the durable context around a durable handler as a convenience to run directly
+        @durable_handler
+        def handler(event: Any, context: DurableContext):  # noqa: ARG001
+            return context_function(*args, **kwargs)(context)
+
+        super().__init__(handler)

tests/checkpoint/processors/base_test.py

@@ -66,9 +66,9 @@ def create_context_details(self, update):
         """Public method to access _create_context_details for testing."""
         return self._create_context_details(update)
 
-    def create_step_details(self, update):
+    def create_step_details(self, update, current_operation):
         """Public method to access _create_step_details for testing."""
-        return self._create_step_details(update)
+        return self._create_step_details(update, current_operation)
 
     def create_callback_details(self, update):
         """Public method to access _create_callback_details for testing."""
@@ -187,7 +187,11 @@ def test_create_step_details():
         error=error,
     )
 
-    result = processor.create_step_details(update)
+    current_op = Mock()
+    current_op.step_details = Mock()
+    current_op.step_details.attempt = Mock()
+
+    result = processor.create_step_details(update, current_op)
 
     assert isinstance(result, StepDetails)
     assert result.result == "test-payload"
@@ -203,11 +207,34 @@ def test_create_step_details_non_step_type():
         payload="test-payload",
     )
 
-    result = processor.create_step_details(update)
+    current_op = Mock()
+    current_op.step_details = Mock()
+    current_op.step_details.attempt = Mock()
+
+    result = processor.create_step_details(update, current_op)
 
     assert result is None
 
 
+def test_create_step_details_without_current_operation():
+    processor = MockProcessor()
+    error = ErrorObject.from_message("test error")
+    update = OperationUpdate(
+        operation_id="test-id",
+        operation_type=OperationType.STEP,
+        action=OperationAction.START,
+        payload="test-payload",
+        error=error,
+    )
+
+    result = processor.create_step_details(update, None)
+
+    assert isinstance(result, StepDetails)
+    assert result.result == "test-payload"
+    assert result.error == error
+    assert result.attempt == 0
+
+
 def test_create_callback_details():
     processor = MockProcessor()
     error = ErrorObject.from_message("test error")

tests/e2e/basic_success_path_test.py

@@ -68,16 +68,16 @@ def function_under_test(event: Any, context: DurableContext) -> list[str]:
         result: DurableFunctionTestResult = runner.run(input="input str", timeout=10)
 
     assert result.status is InvocationStatus.SUCCEEDED
-    assert result.result == '["1 2", "3 4 4 3", "5 6"]'
+    assert result.result == ["1 2", "3 4 4 3", "5 6"]
 
     one_result: StepOperation = result.get_step("one")
-    assert one_result.result == '"1 2"'
+    assert one_result.result == "1 2"
 
     two_result: ContextOperation = result.get_context("two")
-    assert two_result.result == '"3 4 4 3"'
+    assert two_result.result == "3 4 4 3"
 
     three_result: StepOperation = result.get_step("three")
-    assert three_result.result == '"5 6"'
+    assert three_result.result == "5 6"
 
     # currently has the optimization where it's not saving child checkpoints after parent done
     # prob should unpick that for test