chore: code review

Author: bchampp

src/aws_durable_execution_sdk_python_testing/execution.py

@@ -28,6 +28,7 @@
 
 # Import AWS exceptions
 from aws_durable_execution_sdk_python_testing.model import (
+    InvocationCompletedDetails,
     StartDurableExecutionInput,
 )
 from aws_durable_execution_sdk_python_testing.token import (
@@ -60,9 +61,7 @@ def __init__(
         self.start_input: StartDurableExecutionInput = start_input
         self.operations: list[Operation] = operations
         self.updates: list[OperationUpdate] = []
-        self.invocation_completions: list[
-            tuple[float, float, str]
-        ] = []  # (start_ts, end_ts, request_id)
+        self.invocation_completions: list[InvocationCompletedDetails] = []
         self.used_tokens: set[str] = set()
         # TODO: this will need to persist/rehydrate depending on inmemory vs sqllite store
         self._token_sequence: int = 0
@@ -105,8 +104,7 @@ def to_dict(self) -> dict[str, Any]:
             "Operations": [op.to_dict() for op in self.operations],
             "Updates": [update.to_dict() for update in self.updates],
             "InvocationCompletions": [
-                [start, end, req_id]
-                for start, end, req_id in self.invocation_completions
+                completion.to_dict() for completion in self.invocation_completions
             ],
             "UsedTokens": list(self.used_tokens),
             "TokenSequence": self._token_sequence,
@@ -137,7 +135,8 @@ def from_dict(cls, data: dict[str, Any]) -> Execution:
             OperationUpdate.from_dict(update_data) for update_data in data["Updates"]
         ]
         execution.invocation_completions = [
-            tuple(item) for item in data.get("InvocationCompletions", [])
+            InvocationCompletedDetails.from_dict(item)
+            for item in data.get("InvocationCompletions", [])
         ]
         execution.used_tokens = set(data["UsedTokens"])
         execution._token_sequence = data["TokenSequence"]  # noqa: SLF001
@@ -229,7 +228,13 @@ def record_invocation_completion(
         self, start_timestamp: float, end_timestamp: float, request_id: str
     ) -> None:
         """Record an invocation completion event."""
-        self.invocation_completions.append((start_timestamp, end_timestamp, request_id))
+        self.invocation_completions.append(
+            InvocationCompletedDetails(
+                start_timestamp=start_timestamp,
+                end_timestamp=end_timestamp,
+                request_id=request_id,
+            )
+        )
 
     def complete_success(self, result: str | None) -> None:
         """Complete execution successfully (DecisionType.COMPLETE_WORKFLOW_EXECUTION)."""

src/aws_durable_execution_sdk_python_testing/executor.py

@@ -33,6 +33,8 @@
 from aws_durable_execution_sdk_python_testing.model import (
     CheckpointDurableExecutionResponse,
     CheckpointUpdatedExecutionState,
+    EventCreationContext,
+    EventType,
     GetDurableExecutionHistoryResponse,
     GetDurableExecutionResponse,
     GetDurableExecutionStateResponse,
@@ -45,7 +47,6 @@
     StartDurableExecutionOutput,
     StopDurableExecutionResponse,
     TERMINAL_STATUSES,
-    EventCreationContext,
 )
 from aws_durable_execution_sdk_python_testing.model import (
     Event as HistoryEvent,
@@ -415,16 +416,14 @@ def get_execution_history(
         durable_execution_arn: str = execution.durable_execution_arn
 
         # Add InvocationCompleted events
-        for start_ts, end_ts, request_id in execution.invocation_completions:
+        for completion in execution.invocation_completions:
             invocation_event = HistoryEvent(
                 event_id=0,  # Temporary, will be reassigned
-                event_type="InvocationCompleted",
-                event_timestamp=datetime.fromtimestamp(end_ts, tz=UTC),
-                invocation_completed_details={
-                    "StartTimestamp": start_ts,
-                    "EndTimestamp": end_ts,
-                    "RequestId": request_id,
-                },
+                event_type=EventType.INVOCATION_COMPLETED.value,
+                event_timestamp=datetime.fromtimestamp(
+                    completion.end_timestamp, tz=UTC
+                ),
+                invocation_completed_details=completion,
             )
             all_events.append(invocation_event)
 
@@ -785,7 +784,7 @@ async def invoke() -> None:
                 self._store.save(execution)
 
                 invocation_start = time.time()
-                response, request_id = self._invoker.invoke(
+                invoke_response = self._invoker.invoke(
                     execution.start_input.function_name,
                     invocation_input,
                     execution.start_input.lambda_endpoint,
@@ -797,7 +796,7 @@ async def invoke() -> None:
 
                 # Record invocation completion and save immediately
                 execution.record_invocation_completion(
-                    invocation_start, invocation_end, request_id
+                    invocation_start, invocation_end, invoke_response.request_id
                 )
                 self._store.save(execution)
 
@@ -809,6 +808,7 @@ async def invoke() -> None:
                     return
 
                 # Process successful received response - validate status and handle accordingly
+                response = invoke_response.invocation_output
                 try:
                     self._validate_invocation_response_and_store(
                         execution_arn, response, execution

src/aws_durable_execution_sdk_python_testing/invoker.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import json
+from dataclasses import dataclass
 from threading import Lock
 from typing import TYPE_CHECKING, Any, Protocol
 from uuid import uuid4
@@ -27,6 +28,14 @@
     from aws_durable_execution_sdk_python_testing.execution import Execution
 
 
+@dataclass(frozen=True)
+class InvokeResponse:
+    """Response from invoking a durable function."""
+
+    invocation_output: DurableExecutionInvocationOutput
+    request_id: str
+
+
 def create_test_lambda_context() -> LambdaContext:
     # Create client context as a dictionary, not as objects
     # LambdaContext.__init__ expects dictionaries and will create the objects internally
@@ -66,7 +75,7 @@ def invoke(
         function_name: str,
         input: DurableExecutionInvocationInput,
         endpoint_url: str | None = None,
-    ) -> tuple[DurableExecutionInvocationOutput, str]: ...  # pragma: no cover
+    ) -> InvokeResponse: ...  # pragma: no cover
 
     def update_endpoint(
         self, endpoint_url: str, region_name: str
@@ -97,15 +106,17 @@ def invoke(
         function_name: str,  # noqa: ARG002
         input: DurableExecutionInvocationInput,
         endpoint_url: str | None = None,  # noqa: ARG002
-    ) -> tuple[DurableExecutionInvocationOutput, str]:
+    ) -> InvokeResponse:
         # TODO: reasses if function_name will be used in future
         input_with_client = DurableExecutionInvocationInputWithClient.from_durable_execution_invocation_input(
             input, self.service_client
         )
         context = create_test_lambda_context()
         response_dict = self.handler(input_with_client, context)
         output = DurableExecutionInvocationOutput.from_dict(response_dict)
-        return output, context.aws_request_id
+        return InvokeResponse(
+            invocation_output=output, request_id=context.aws_request_id
+        )
 
     def update_endpoint(self, endpoint_url: str, region_name: str) -> None:
         """No-op for in-process invoker."""
@@ -194,7 +205,7 @@ def invoke(
         function_name: str,
         input: DurableExecutionInvocationInput,
         endpoint_url: str | None = None,
-    ) -> tuple[DurableExecutionInvocationOutput, str]:
+    ) -> InvokeResponse:
         """Invoke AWS Lambda function and return durable execution result.
 
         Args:
@@ -203,7 +214,7 @@ def invoke(
             endpoint_url: Lambda endpoint url
 
         Returns:
-            tuple: (DurableExecutionInvocationOutput, request_id)
+            InvokeResponse: Response containing invocation output and request ID
 
         Raises:
             ResourceNotFoundException: If function does not exist
@@ -259,7 +270,7 @@ def invoke(
 
             # Convert to DurableExecutionInvocationOutput
             output = DurableExecutionInvocationOutput.from_dict(response_dict)
-            return output, request_id
+            return InvokeResponse(invocation_output=output, request_id=request_id)
 
         except client.exceptions.ResourceNotFoundException as e:
             msg = f"Function not found: {function_name}"

src/aws_durable_execution_sdk_python_testing/model.py

@@ -68,6 +68,7 @@ class EventType(Enum):
     CALLBACK_SUCCEEDED = "CallbackSucceeded"
     CALLBACK_FAILED = "CallbackFailed"
     CALLBACK_TIMED_OUT = "CallbackTimedOut"
+    INVOCATION_COMPLETED = "InvocationCompleted"
 
 
 TERMINAL_STATUSES: set[OperationStatus] = {
@@ -1222,6 +1223,30 @@ def to_dict(self) -> dict[str, Any]:
         return result
 
 
+@dataclass(frozen=True)
+class InvocationCompletedDetails:
+    """Invocation completed event details."""
+
+    start_timestamp: float
+    end_timestamp: float
+    request_id: str
+
+    @classmethod
+    def from_dict(cls, data: dict) -> InvocationCompletedDetails:
+        return cls(
+            start_timestamp=data["StartTimestamp"],
+            end_timestamp=data["EndTimestamp"],
+            request_id=data["RequestId"],
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "StartTimestamp": self.start_timestamp,
+            "EndTimestamp": self.end_timestamp,
+            "RequestId": self.request_id,
+        }
+
+
 # endregion event_structures
 
 
@@ -1329,7 +1354,7 @@ class Event:
     callback_succeeded_details: CallbackSucceededDetails | None = None
     callback_failed_details: CallbackFailedDetails | None = None
     callback_timed_out_details: CallbackTimedOutDetails | None = None
-    invocation_completed_details: dict[str, Any] | None = None
+    invocation_completed_details: InvocationCompletedDetails | None = None
 
     @classmethod
     def from_dict(cls, data: dict) -> Event:
@@ -1448,7 +1473,11 @@ def from_dict(cls, data: dict) -> Event:
         if details_data := data.get("CallbackTimedOutDetails"):
             callback_timed_out_details = CallbackTimedOutDetails.from_dict(details_data)
 
-        invocation_completed_details = data.get("InvocationCompletedDetails")
+        invocation_completed_details = None
+        if details_data := data.get("InvocationCompletedDetails"):
+            invocation_completed_details = InvocationCompletedDetails.from_dict(
+                details_data
+            )
 
         return cls(
             event_type=data["EventType"],
@@ -1568,7 +1597,9 @@ def to_dict(self) -> dict[str, Any]:
                 self.callback_timed_out_details.to_dict()
             )
         if self.invocation_completed_details is not None:
-            result["InvocationCompletedDetails"] = self.invocation_completed_details
+            result["InvocationCompletedDetails"] = (
+                self.invocation_completed_details.to_dict()
+            )
         return result
 
     # region execution
@@ -2224,6 +2255,30 @@ def create_callback_event(cls, context: EventCreationContext) -> Event:
 
     # endregion callback
 
+    # region invocation_completed
+    @classmethod
+    def create_invocation_completed(
+        cls,
+        event_id: int,
+        event_timestamp: datetime.datetime,
+        start_timestamp: float,
+        end_timestamp: float,
+        request_id: str,
+    ) -> Event:
+        """Create invocation completed event."""
+        return cls(
+            event_type=EventType.INVOCATION_COMPLETED.value,
+            event_timestamp=event_timestamp,
+            event_id=event_id,
+            invocation_completed_details=InvocationCompletedDetails(
+                start_timestamp=start_timestamp,
+                end_timestamp=end_timestamp,
+                request_id=request_id,
+            ),
+        )
+
+    # endregion invocation_completed
+
     @classmethod
     def create_event_started(cls, context: EventCreationContext) -> Event:
         """Convert operation to started event."""