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_result_path_only_state_input(dummy_resource):
task_state = TaskState("Task", resource=dummy_resource, result_path=None)
state_machine = StateMachine(start_state=task_state)
# Check the output from compiling
assert state_machine.compile() == {
"StartAt": task_state.name,
"States": {
task_state.name: {
"Resource": dummy_resource,
"ResultPath": None,
"Type": "Task",
"End": True,
},
},
}
# Simulate the state machine
state_input = {
"comment": "This is a test of the input and output of a Task state.",
"details": "Default example",
"who": "AWS Step Functions",
}
def mock_fn(event, context):
return "Hello, AWS Step Functions!"
state_output = state_machine.simulate(
state_input,
resource_to_mock_fn={dummy_resource: mock_fn},
)
# Keeps the only the state output
assert state_output == state_input
def test_task_state(dummy_resource, capture_stdout):
pass_state = PassState("Pass", comment="The starting state")
timeout_seconds = 10
task_state = TaskState(
"Task", resource=dummy_resource, timeout_seconds=timeout_seconds
)
# Define the state machine
pass_state >> task_state
state_machine = StateMachine(start_state=pass_state)
# Check the output from compiling
assert state_machine.compile() == {
"StartAt": pass_state.name,
"States": {
pass_state.name: {
"Comment": pass_state.comment,
"Type": "Pass",
"Next": task_state.name,
},
task_state.name: {
"Type": "Task",
"Resource": dummy_resource,
"TimeoutSeconds": timeout_seconds,
"End": True,
},
},
}
# Simulate the state machine
def mock_fn(event, context):
event["foo"] *= 2
return event
stdout = capture_stdout(
lambda: state_machine.simulate(
{"foo": 5, "bar": 1},
resource_to_mock_fn={dummy_resource: mock_fn},
)
)
assert (
stdout
== """Starting simulation of state machine
Executing PassState('Pass')
State input: {'foo': 5, 'bar': 1}
State input after applying input path of $: {'foo': 5, 'bar': 1}
Output from applying result path of $: {'foo': 5, 'bar': 1}
State output after applying output path of $: {'foo': 5, 'bar': 1}
State output: {'foo': 5, 'bar': 1}
Executing TaskState('Task')
State input: {'foo': 5, 'bar': 1}
State input after applying input path of $: {'foo': 5, 'bar': 1}
Output from applying result path of $: {'foo': 10, 'bar': 1}
State output after applying output path of $: {'foo': 10, 'bar': 1}
State output: {'foo': 10, 'bar': 1}
Terminating simulation of state machine
"""
)
def test_state_has_invalid_result_selector(dummy_resource):
invalid_result_selector = {"ClusterId.$": "$.dataset*"}
with pytest.raises(
AWSStepFuncsValueError, match='Unsupported Reference Path operator: "*"'
):
TaskState(
"My Task",
resource=dummy_resource,
result_selector=invalid_result_selector,
)
def test_compile_parameters():
y_state = PassState("Y")
x_state = TaskState(
"X",
resource="arn:aws:states:us-east-1:123456789012:task:X",
parameters={
"first": 88,
"second": 99
},
)
_ = x_state >> y_state
assert x_state.compile() == {
"Type": "Task",
"Resource": "arn:aws:states:us-east-1:123456789012:task:X",
"Next": "Y",
"Parameters": {
"first": 88,
"second": 99
},
}
def test_result_selector(dummy_resource):
result_selector = {
"ClusterId.$": "$.output.ClusterId",
"ResourceType.$": "$.resourceType",
"SomethingElse.$": "$.keyDoesntExist",
}
task_state = TaskState(
"Task", resource=dummy_resource, result_selector=result_selector
)
state_machine = StateMachine(start_state=task_state)
# Check the output from compiling
assert state_machine.compile() == {
"StartAt": task_state.name,
"States": {
task_state.name: {
"Resource": dummy_resource,
"ResultSelector": result_selector,
"Type": "Task",
"End": True,
},
},
}
# Simulate the state machine
def mock_fn(event, context):
return {
"resourceType": "elasticmapreduce",
"resource": "createCluster.sync",
"output": {
"SdkHttpMetadata": {
"HttpHeaders": {
"Content-Length": "1112",
"Content-Type": "application/x-amz-JSON-1.1",
"Date": "Mon, 25 Nov 2019 19:41:29 GMT",
"x-amzn-RequestId": "1234-5678-9012",
},
"HttpStatusCode": 200,
},
"SdkResponseMetadata": {"RequestId": "1234-5678-9012"},
"ClusterId": "AKIAIOSFODNN7EXAMPLE",
},
}
state_output = state_machine.simulate(
resource_to_mock_fn={dummy_resource: mock_fn},
)
assert state_output == {
"ResourceType": "elasticmapreduce",
"ClusterId": "AKIAIOSFODNN7EXAMPLE",
}
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_retrier_zero_max_attempts():
task_state = TaskState("Task", resource="123").add_retrier(
["SomeError"], max_attempts=0
)
fail_state = FailState("Fail", error="SomeError", cause="I did it!")
task_state >> fail_state
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
"""
)
def iterator(resource):
task_state = TaskState("Validate", resource=resource)
return StateMachine(start_state=task_state)