chore: fix mocks for start input to have execution timeout

Author: bchampp

src/aws_durable_execution_sdk_python_testing/executor.py

@@ -120,24 +120,19 @@ def start_execution(
         self._completion_events[execution.durable_execution_arn] = completion_event
 
         # Schedule execution timeout
-        try:
-            timeout_seconds = input.execution_timeout_seconds
-            if timeout_seconds and timeout_seconds > 0:
-
-                def timeout_handler():
-                    error = ErrorObject.from_message(
-                        f"Execution timed out after {timeout_seconds} seconds."
-                    )
-                    self.on_timed_out(execution.durable_execution_arn, error)
+        if input.execution_timeout_seconds > 0:
 
-                self._execution_timeout = self._scheduler.call_later(
-                    timeout_handler,
-                    delay=timeout_seconds,
-                    completion_event=completion_event,
+            def timeout_handler():
+                error = ErrorObject.from_message(
+                    f"Execution timed out after {input.execution_timeout_seconds} seconds."
                 )
-        except (AttributeError, TypeError):
-            # Handle Mock objects or invalid timeout values in tests
-            pass
+                self.on_timed_out(execution.durable_execution_arn, error)
+
+            self._execution_timeout = self._scheduler.call_later(
+                timeout_handler,
+                delay=input.execution_timeout_seconds,
+                completion_event=completion_event,
+            )
 
         # Schedule initial invocation to run immediately
         self._invoke_execution(execution.durable_execution_arn)

tests/executor_test.py

@@ -1213,6 +1213,7 @@ def test_wait_until_complete_success(executor, mock_scheduler):
         mock_execution_class.new.return_value = mock_execution
 
         start_input = Mock()
+        start_input.execution_timeout_seconds = 0
         executor.start_execution(start_input)
 
     result = executor.wait_until_complete("test-arn", timeout=10)
@@ -1236,6 +1237,7 @@ def test_wait_until_complete_timeout(executor, mock_scheduler):
         mock_execution_class.new.return_value = mock_execution
 
         start_input = Mock()
+        start_input.execution_timeout_seconds = 0
         executor.start_execution(start_input)
 
     result = executor.wait_until_complete("test-arn", timeout=10)
@@ -1290,6 +1292,7 @@ def test_should_schedule_wait_timer_correctly(executor, mock_scheduler):
         mock_execution_class.new.return_value = mock_execution
 
         start_input = Mock()
+        start_input.execution_timeout_seconds = 0
         executor.start_execution(start_input)
 
     # Act - schedule wait timer through public method
@@ -1444,6 +1447,7 @@ def test_on_wait_timer_scheduled(executor, mock_scheduler):
         mock_execution_class.new.return_value = mock_execution
 
         start_input = Mock()
+        start_input.execution_timeout_seconds = 0
         executor.start_execution(start_input)
 
     with patch.object(executor, "_on_wait_succeeded"):
@@ -1654,6 +1658,7 @@ def test_invoke_execution_with_delay_through_wait_timer(executor, mock_scheduler
         mock_execution_class.new.return_value = mock_execution
 
         start_input = Mock()
+        start_input.execution_timeout_seconds = 0
         executor.start_execution(start_input)
 
     # Test delay behavior through wait timer scheduling
@@ -1681,6 +1686,7 @@ def test_invoke_execution_no_delay_through_start_execution(executor, mock_schedu
         mock_execution_class.new.return_value = mock_execution
 
         start_input = Mock()
+        start_input.execution_timeout_seconds = 0
         executor.start_execution(start_input)
 
     # Verify scheduler was called with no delay for initial execution
@@ -1704,6 +1710,7 @@ def test_on_step_retry_scheduled(executor, mock_scheduler):
         mock_execution_class.new.return_value = mock_execution
 
         start_input = Mock()
+        start_input.execution_timeout_seconds = 0
         executor.start_execution(start_input)
 
     with patch.object(executor, "_on_retry_ready"):
@@ -1733,6 +1740,7 @@ def test_wait_handler_execution(executor, mock_scheduler):
         mock_execution_class.new.return_value = mock_execution
 
         start_input = Mock()
+        start_input.execution_timeout_seconds = 0
         executor.start_execution(start_input)
 
     with patch.object(executor, "_on_wait_succeeded") as mock_wait:
@@ -1764,6 +1772,7 @@ def test_retry_handler_execution(executor, mock_scheduler):
         mock_execution_class.new.return_value = mock_execution
 
         start_input = Mock()
+        start_input.execution_timeout_seconds = 0
         executor.start_execution(start_input)
 
     with patch.object(executor, "_on_retry_ready") as mock_retry: