def test_multiple_catchers(capture_stdout): resource = "123" task_state = TaskState("Task", resource=resource) timeout_state = PassState("Timeout") task_failed_state = PassState("Task Failed") task_state.add_catcher(["States.Timeout"], next_state=timeout_state).add_catcher( ["States.TaskFailed"], next_state=task_failed_state ) state_machine = StateMachine(start_state=task_state) assert state_machine.compile() == { "StartAt": "Task", "States": { "Task Failed": {"Type": "Pass", "End": True}, "Timeout": {"Type": "Pass", "End": True}, "Task": { "Type": "Task", "End": True, "Catch": [ {"ErrorEquals": ["States.Timeout"], "Next": "Timeout"}, {"ErrorEquals": ["States.TaskFailed"], "Next": "Task Failed"}, ], "Resource": "123", }, }, } def failure_mock_fn(event, context): # Will cause a TaskFailedError assert False # noqa: PT015 stdout = capture_stdout( lambda: state_machine.simulate(resource_to_mock_fn={resource: failure_mock_fn}) ) assert ( stdout == """Starting simulation of state machine Executing TaskState('Task') State input: {} State input after applying input path of $: {} TaskFailedError encountered in state Checking for catchers Found catcher, transitioning to PassState('Task Failed') State output: {} Executing PassState('Task Failed') State input: {} State input after applying input path of $: {} Output from applying result path of $: {} State output after applying output path of $: {} State output: {} Terminating simulation of state machine """ )
def test_duplicate_names_catcher(): duplicate_name = "Duplicated" task_state = TaskState(duplicate_name, resource="123") transition_state = FailState(duplicate_name, error="MyError", cause="Negligence") task_state.add_catcher(["Something"], next_state=transition_state) with pytest.raises( AWSStepFuncsValueError, match= "Duplicate names detected in state machine. Names must be unique", ): StateMachine(start_state=task_state)
def test_compile_retrier_and_catcher(): task_state = TaskState("Task", resource="123").add_retrier( ["SomeError"], max_attempts=0 ) transition_state = TaskState("Cleanup", resource="456") task_state.add_catcher(["States.ALL"], next_state=transition_state) assert task_state.compile() == { "Type": "Task", "End": True, "Retry": [{"ErrorEquals": ["SomeError"], "MaxAttempts": 0}], "Catch": [{"ErrorEquals": ["States.ALL"], "Next": "Cleanup"}], "Resource": "123", }
def test_catcher(capture_stdout): resource = "123" task_state = TaskState("Task", resource=resource) succeed_state = SucceedState("Success") pass_state = PassState("Pass") fail_state = FailState("Failure", error="IFailed", cause="I failed!") task_state >> succeed_state pass_state >> fail_state task_state.add_catcher(["States.ALL"], next_state=pass_state) state_machine = StateMachine(start_state=task_state) def failure_mock_fn(event, context): assert False # noqa: PT015 stdout = capture_stdout( lambda: state_machine.simulate(resource_to_mock_fn={resource: failure_mock_fn}) ) assert ( stdout == """Starting simulation of state machine Executing TaskState('Task') State input: {} State input after applying input path of $: {} TaskFailedError encountered in state Checking for catchers Found catcher, transitioning to PassState('Pass') State output: {} Executing PassState('Pass') State input: {} State input after applying input path of $: {} Output from applying result path of $: {} State output after applying output path of $: {} State output: {} Executing FailState('Failure', error='IFailed', cause='I failed!') State input: {} FailStateError encountered in state Checking for catchers State output: {} Terminating simulation of state machine """ ) # Test no catcher matched task_state.catchers.pop() # Remove States.ALL catcher task_state.add_catcher(["Timeout"], next_state=pass_state) stdout = capture_stdout( lambda: state_machine.simulate(resource_to_mock_fn={resource: failure_mock_fn}) ) assert ( stdout == """Starting simulation of state machine Executing TaskState('Task') State input: {} State input after applying input path of $: {} TaskFailedError encountered in state Checking for catchers No catchers were matched State output: {} Terminating simulation of state machine """ )