diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/__init__.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/__init__.py index ab62838c..5691e92b 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/__init__.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/__init__.py @@ -21,6 +21,7 @@ DurableExecutionsError, InvocationError, ValidationError, + WaitForConditionError, ) # Core decorator - used in every durable function @@ -39,6 +40,7 @@ "ParallelBranch", "StepContext", "ValidationError", + "WaitForConditionError", "WithRetryConfig", "__version__", "durable_execution", diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/exceptions.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/exceptions.py index a3fe3196..30abff6d 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/exceptions.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/exceptions.py @@ -266,6 +266,12 @@ class InvalidStateError(DurableExecutionsError): """Raised when an operation is attempted on an object in an invalid state.""" +class WaitForConditionError(DurableExecutionsError): + """Raised when a wait_for_condition exhausts its polling attempts before the + condition is met, matching the exhaustion semantics of the JS and Java SDKs. + """ + + class UserlandError(DurableExecutionsError): """Failure in user-land - i.e code passed into durable executions from the caller.""" diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/wait_for_condition.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/wait_for_condition.py index 8771de36..0e07693d 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/wait_for_condition.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/operation/wait_for_condition.py @@ -239,7 +239,7 @@ def execute(self, checkpointed_result: CheckpointedResult) -> T: if delay_seconds is not None and delay_seconds < 1: logger.warning( ( - "WaitDecision delay_seconds step for id: %s, name: %s," + "wait_for_condition delay_seconds for id: %s, name: %s," "is %d < 1. Setting to minimum of 1 seconds." ), self.operation_identifier.operation_id, diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/waits.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/waits.py index b4d740a6..e5aa325d 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/waits.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/waits.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Generic from aws_durable_execution_sdk_python.config import Duration, JitterStrategy, T +from aws_durable_execution_sdk_python.exceptions import WaitForConditionError if TYPE_CHECKING: from collections.abc import Callable @@ -16,29 +17,6 @@ Numeric = int | float -@dataclass -class WaitDecision: - """Decision about whether to wait a step and with what delay.""" - - should_wait: bool - delay: Duration - - @property - def delay_seconds(self) -> int: - """Get delay in seconds.""" - return self.delay.to_seconds() - - @classmethod - def wait(cls, delay: Duration) -> WaitDecision: - """Create a wait decision.""" - return cls(should_wait=True, delay=delay) - - @classmethod - def no_wait(cls) -> WaitDecision: - """Create a no-wait decision.""" - return cls(should_wait=False, delay=Duration()) - - @dataclass class WaitStrategyConfig(Generic[T]): should_continue_polling: Callable[[T], bool] @@ -69,35 +47,6 @@ def timeout_seconds(self) -> int | None: return self.timeout.to_seconds() -def create_wait_strategy( - config: WaitStrategyConfig[T], -) -> Callable[[T, int], WaitDecision]: - def wait_strategy(result: T, attempts_made: int) -> WaitDecision: - # Check if condition is met - if not config.should_continue_polling(result): - return WaitDecision.no_wait() - - # Check if we've exceeded max attempts - if attempts_made >= config.max_attempts: - return WaitDecision.no_wait() - - # Calculate delay with exponential backoff - base_delay: float = min( - config.initial_delay_seconds * (config.backoff_rate ** (attempts_made - 1)), - config.max_delay_seconds, - ) - - # Apply jitter to get final delay - delay_with_jitter: float = config.jitter_strategy.apply_jitter(base_delay) - - # Round up and ensure minimum of 1 second - final_delay: int = max(1, math.ceil(delay_with_jitter)) - - return WaitDecision.wait(Duration(seconds=final_delay)) - - return wait_strategy - - @dataclass(frozen=True) class WaitForConditionDecision: """Decision about whether to continue waiting.""" @@ -121,6 +70,41 @@ def stop_polling(cls) -> WaitForConditionDecision: return cls(should_continue=False, delay=Duration()) +def create_wait_strategy( + config: WaitStrategyConfig[T], +) -> Callable[[T, int], WaitForConditionDecision]: + def wait_strategy(result: T, attempts_made: int) -> WaitForConditionDecision: + # Condition satisfied wins over exhaustion, so a condition met on the final + # attempt still succeeds (matches the JS and Java SDKs). + if not config.should_continue_polling(result): + return WaitForConditionDecision.stop_polling() + + # Out of attempts: fail rather than stop, otherwise the executor would treat + # this as success and return partial state. + if attempts_made >= config.max_attempts: + msg = ( + f"wait_for_condition exhausted {config.max_attempts} attempts " + "before the condition was met" + ) + raise WaitForConditionError(msg) + + # Calculate delay with exponential backoff + base_delay: float = min( + config.initial_delay_seconds * (config.backoff_rate ** (attempts_made - 1)), + config.max_delay_seconds, + ) + + # Apply jitter to get final delay + delay_with_jitter: float = config.jitter_strategy.apply_jitter(base_delay) + + # Round up and ensure minimum of 1 second + final_delay: int = max(1, math.ceil(delay_with_jitter)) + + return WaitForConditionDecision.continue_waiting(Duration(seconds=final_delay)) + + return wait_strategy + + @dataclass(frozen=True) class WaitForConditionConfig(Generic[T]): """Configuration for wait_for_condition.""" diff --git a/packages/aws-durable-execution-sdk-python/tests/operation/wait_for_condition_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/wait_for_condition_test.py index da35a6fe..bf729fb4 100644 --- a/packages/aws-durable-execution-sdk-python/tests/operation/wait_for_condition_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/operation/wait_for_condition_test.py @@ -11,6 +11,7 @@ CallableRuntimeError, InvocationError, SuspendExecution, + WaitForConditionError, ) from aws_durable_execution_sdk_python.identifier import OperationIdentifier from aws_durable_execution_sdk_python.lambda_service import ( @@ -30,6 +31,8 @@ from aws_durable_execution_sdk_python.waits import ( WaitForConditionConfig, WaitForConditionDecision, + WaitStrategyConfig, + create_wait_strategy, ) from tests.serdes_test import CustomDictSerDes @@ -1386,6 +1389,95 @@ def mock_wait_strategy(state, attempt): assert mock_state.create_checkpoint.call_count == 2 # START + SUCCESS checkpoints +def test_wait_for_condition_exhaustion_raises_and_checkpoints_fail(): + """Live path: the built-in strategy runs out of attempts, so it raises + WaitForConditionError, which is checkpointed as a FAIL and propagated.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "arn:aws:test" + mock_state.get_checkpoint_result.return_value = ( + CheckpointedResult.create_not_found() + ) + + mock_logger = Mock(spec=Logger) + mock_logger.with_log_info.return_value = mock_logger + + op_id = OperationIdentifier( + "op1", OperationSubType.WAIT_FOR_CONDITION, None, "test_wait" + ) + + def check_func(state, context): + return state + 1 + + mock_state.wrap_user_function.return_value = check_func + + # max_attempts=1 means attempt 1 is already the last one. + config = WaitForConditionConfig( + initial_state=5, + wait_strategy=create_wait_strategy( + WaitStrategyConfig(should_continue_polling=lambda x: True, max_attempts=1) + ), + ) + + with pytest.raises(WaitForConditionError): + wait_for_condition_handler( + state=mock_state, + operation_identifier=op_id, + check=check_func, + config=config, + context_logger=mock_logger, + ) + + assert mock_state.create_checkpoint.call_count == 2 # START and FAIL + fail_operation = mock_state.create_checkpoint.call_args_list[1][1][ + "operation_update" + ] + assert fail_operation.error.type == "WaitForConditionError" + + +def test_wait_for_condition_exhaustion_surfaces_on_replay(): + """Replay path: the FAILED checkpoint short-circuits on the next invocation. + Typed reconstruction lands with #120, so today it surfaces as a + CallableRuntimeError carrying the original error_type.""" + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + operation = Operation( + operation_id="op1", + operation_type=OperationType.STEP, + status=OperationStatus.FAILED, + step_details=StepDetails( + error=ErrorObject("exhausted attempts", "WaitForConditionError", None, None) + ), + ) + mock_result = CheckpointedResult.create_from_operation(operation) + mock_state.get_checkpoint_result.return_value = mock_result + + mock_logger = Mock(spec=Logger) + op_id = OperationIdentifier( + "op1", OperationSubType.WAIT_FOR_CONDITION, None, "test_wait" + ) + + def check_func(state, context): + msg = "Check function should not be called on replay of a failure" + raise AssertionError(msg) + + config = WaitForConditionConfig( + initial_state=5, + wait_strategy=lambda s, a: WaitForConditionDecision.stop_polling(), + ) + + with pytest.raises(CallableRuntimeError) as exc_info: + wait_for_condition_handler( + state=mock_state, + operation_identifier=op_id, + check=check_func, + config=config, + context_logger=mock_logger, + ) + + assert exc_info.value.error_type == "WaitForConditionError" + assert mock_state.create_checkpoint.call_count == 0 # Nothing new on replay + + def test_wait_for_condition_executes_check_when_checkpoint_not_terminal_duplicate(): """Test backward compatibility: when checkpoint is not terminal (STARTED), the wait_for_condition operation executes the check function normally. diff --git a/packages/aws-durable-execution-sdk-python/tests/waits_test.py b/packages/aws-durable-execution-sdk-python/tests/waits_test.py index 06267d8d..47c050dc 100644 --- a/packages/aws-durable-execution-sdk-python/tests/waits_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/waits_test.py @@ -2,10 +2,12 @@ from unittest.mock import patch +import pytest + from aws_durable_execution_sdk_python.config import Duration, JitterStrategy +from aws_durable_execution_sdk_python.exceptions import WaitForConditionError from aws_durable_execution_sdk_python.serdes import JsonSerDes from aws_durable_execution_sdk_python.waits import ( - WaitDecision, WaitForConditionConfig, WaitForConditionDecision, WaitStrategyConfig, @@ -13,22 +15,6 @@ ) -class TestWaitDecision: - """Test WaitDecision factory methods.""" - - def test_wait_factory(self): - """Test wait factory method.""" - decision = WaitDecision.wait(Duration.from_seconds(30)) - assert decision.should_wait is True - assert decision.delay_seconds == 30 - - def test_no_wait_factory(self): - """Test no_wait factory method.""" - decision = WaitDecision.no_wait() - assert decision.should_wait is False - assert decision.delay_seconds == 0 - - class TestWaitForConditionDecision: """Test WaitForConditionDecision factory methods.""" @@ -62,25 +48,34 @@ def test_default_config(self): class TestCreateWaitStrategy: """Test wait strategy creation and behavior.""" - def test_condition_met_returns_no_wait(self): - """Test strategy returns no_wait when condition is met.""" + def test_condition_met_stops_polling(self): + """Test strategy stops polling when the condition is met.""" config = WaitStrategyConfig(should_continue_polling=lambda x: False) strategy = create_wait_strategy(config) result = "completed" decision = strategy(result, 1) - assert decision.should_wait is False + assert decision.should_continue is False - def test_max_attempts_exceeded(self): - """Test strategy returns no_wait when max attempts exceeded.""" + def test_max_attempts_exceeded_raises(self): + """Test strategy raises once attempts are exhausted.""" config = WaitStrategyConfig( should_continue_polling=lambda x: True, max_attempts=5 ) strategy = create_wait_strategy(config) - result = "pending" - decision = strategy(result, 5) - assert decision.should_wait is False + with pytest.raises(WaitForConditionError): + strategy("pending", 5) + + def test_condition_met_on_final_attempt_still_succeeds(self): + """Condition satisfied on the last attempt wins over exhaustion.""" + config = WaitStrategyConfig( + should_continue_polling=lambda x: False, max_attempts=5 + ) + strategy = create_wait_strategy(config) + + decision = strategy("done", 5) + assert decision.should_continue is False def test_should_continue_polling(self): """Test strategy continues when condition not met.""" @@ -89,7 +84,7 @@ def test_should_continue_polling(self): result = "pending" decision = strategy(result, 1) - assert decision.should_wait is True + assert decision.should_continue is True @patch("random.random") def test_exponential_backoff_calculation(self, mock_random): @@ -205,12 +200,12 @@ def __init__(self, count): # Should continue when count < 3 state1 = State(1) decision1 = strategy(state1, 1) - assert decision1.should_wait is True + assert decision1.should_continue is True # Should stop when count >= 3 state2 = State(3) decision2 = strategy(state2, 1) - assert decision2.should_wait is False + assert decision2.should_continue is False def test_complex_condition_logic(self): """Test complex condition logic.""" @@ -224,17 +219,17 @@ def complex_condition(result): # Should continue result1 = {"status": "pending", "retries": 2} decision1 = strategy(result1, 1) - assert decision1.should_wait is True + assert decision1.should_continue is True # Should stop - status changed result2 = {"status": "completed", "retries": 2} decision2 = strategy(result2, 1) - assert decision2.should_wait is False + assert decision2.should_continue is False # Should stop - retries exceeded result3 = {"status": "pending", "retries": 5} decision3 = strategy(result3, 1) - assert decision3.should_wait is False + assert decision3.should_continue is False class TestEdgeCases: @@ -295,13 +290,13 @@ def test_attempt_at_boundary(self): result = "pending" - # At boundary - should not wait - decision = strategy(result, 3) - assert decision.should_wait is False + # At boundary - out of attempts, so it fails + with pytest.raises(WaitForConditionError): + strategy(result, 3) - # Just before boundary - should wait + # Just before boundary - should continue decision = strategy(result, 2) - assert decision.should_wait is True + assert decision.should_continue is True def test_negative_delay_clamped_to_one(self): """Test negative delay is clamped to 1.""" @@ -364,6 +359,18 @@ def wait_strategy(state, attempt): assert config.serdes is serdes + def test_built_in_strategy_wires_into_config(self): + """The built-in strategy returns the type wait_for_condition consumes.""" + config = WaitForConditionConfig( + wait_strategy=create_wait_strategy( + WaitStrategyConfig(should_continue_polling=lambda x: x < 3) + ), + initial_state=0, + ) + + decision = config.wait_strategy(1, 1) + assert decision.should_continue is True + class TestWaitStrategyCallableConditions: """Test wait strategy with various callable conditions.""" @@ -374,10 +381,10 @@ def test_lambda_condition(self): strategy = create_wait_strategy(config) decision1 = strategy(5, 1) - assert decision1.should_wait is True + assert decision1.should_continue is True decision2 = strategy(10, 1) - assert decision2.should_wait is False + assert decision2.should_continue is False def test_function_condition(self): """Test with function condition.""" @@ -389,10 +396,10 @@ def is_pending(status): strategy = create_wait_strategy(config) decision1 = strategy("pending", 1) - assert decision1.should_wait is True + assert decision1.should_continue is True decision2 = strategy("completed", 1) - assert decision2.should_wait is False + assert decision2.should_continue is False def test_method_condition(self): """Test with method condition.""" @@ -409,7 +416,7 @@ def should_continue(self, value): strategy = create_wait_strategy(config) decision1 = strategy(50, 1) - assert decision1.should_wait is True + assert decision1.should_continue is True decision2 = strategy(100, 1) - assert decision2.should_wait is False + assert decision2.should_continue is False