def test_pass_state(capture_stdout):
pass_state1 = PassState("Pass 1", comment="The starting state")
pass_state2 = PassState("Pass 2")
pass_state3 = PassState("Pass 3")
pass_state1 >> pass_state2 >> pass_state3
state_machine = StateMachine(start_state=pass_state1)
assert [state.name for state in state_machine.start_state] == [
"Pass 1",
"Pass 2",
"Pass 3",
]
assert state_machine.compile() == {
"StartAt": "Pass 1",
"States": {
"Pass 2": {
"Type": "Pass",
"Next": "Pass 3"
},
"Pass 1": {
"Type": "Pass",
"Comment": "The starting state",
"Next": "Pass 2",
},
"Pass 3": {
"Type": "Pass",
"End": True
},
},
}
stdout = capture_stdout(lambda: state_machine.simulate())
assert (stdout == """Starting simulation of state machine
Executing PassState('Pass 1')
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 PassState('Pass 2')
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 PassState('Pass 3')
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():
duplicate_name = "My Pass"
pass_state1 = PassState(duplicate_name)
pass_state2 = PassState(duplicate_name)
pass_state1 >> pass_state2
with pytest.raises(
AWSStepFuncsValueError,
match=
"Duplicate names detected in state machine. Names must be unique",
):
StateMachine(start_state=pass_state1)
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_pass_state_result(capture_stdout):
result = {"Hello": "world!"}
pass_state = PassState("Passing", result=result)
state_machine = StateMachine(start_state=pass_state)
assert state_machine.compile() == {
"StartAt": "Passing",
"States": {
"Passing": {
"Type": "Pass",
"End": True,
"Result": {
"Hello": "world!"
}
}
},
}
stdout = capture_stdout(lambda: state_machine.simulate())
assert (stdout == """Starting simulation of state machine
Executing PassState('Passing')
State input: {}
State input after applying input path of $: {}
Output from applying result path of $: {'Hello': 'world!'}
State output after applying output path of $: {'Hello': 'world!'}
State output: {'Hello': 'world!'}
Terminating simulation of state machine
""")
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_terminal_state():
fail_state = FailState("Fail",
error="JustBecause",
cause="Because I feel like it")
pass_state = PassState("Pass")
with pytest.raises(AWSStepFuncsValueError,
match="FailState cannot have a next state"):
fail_state >> pass_state
def test_state_transition_reprs():
pass_state = PassState("Pass", comment="This is a pass state")
fail_state = FailState("Fail", error="MyBad", cause="It's a problem")
pass_state >> fail_state
assert (
repr(pass_state) ==
"PassState(name='Pass', comment='This is a pass state', next_state='Fail')"
)
assert repr(fail_state) == "FailState(name='Fail')"
def test_not_choice():
next_state = PassState("Passing")
not_choice = NotChoice(
variable="$.type",
string_equals="Private",
next_state=next_state,
)
assert not_choice.evaluate({"type": "Public"})
assert not not_choice.evaluate({"type": "Private"})
assert not_choice.evaluate({"sex": "Male"})
def test_and_choice():
next_state = PassState("Passing")
and_choice = AndChoice(
[
ChoiceRule(variable="$.value", is_present=True),
ChoiceRule(variable="$.value", numeric_greater_than_equals=20),
ChoiceRule(variable="$.value", numeric_less_than=30),
],
next_state=next_state,
)
assert and_choice.evaluate({"setting": "on", "value": 20})
assert and_choice.evaluate({"setting": "on", "value": 25})
assert not and_choice.evaluate({"setting": "on", "value": 30})
assert not and_choice.evaluate({"setting": "on"})
assert not and_choice.evaluate({"setting": "on", "value": 50})
def test_variable_choice():
next_state = PassState("Passing")
variable_choice = VariableChoice(
variable="$.type",
string_equals="Private",
next_state=next_state,
)
assert not variable_choice.evaluate({"type": "Public"})
assert variable_choice.evaluate({"type": "Private"})
variable_choice = VariableChoice(
variable="$.rating",
numeric_greater_than_path="$.auditThreshold",
next_state=next_state,
)
assert not variable_choice.evaluate({"rating": 53, "auditThreshold": 60})
assert variable_choice.evaluate({"rating": 53, "auditThreshold": 50})
assert not variable_choice.evaluate({"rating": 53, "auditThreshold": 53})
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_path():
pass_state = PassState(
"Passing", result={"Hello": "world!"}, result_path="$.result"
)
state_machine = StateMachine(start_state=pass_state)
assert state_machine.compile() == {
"StartAt": "Passing",
"States": {
"Passing": {
"Type": "Pass",
"End": True,
"ResultPath": "$.result",
"Result": {"Hello": "world!"},
}
},
}
state_output = state_machine.simulate({"sum": 42})
assert state_output == {"sum": 42, "result": {"Hello": "world!"}}
def test_to_json(tmp_path):
pass_state = PassState("My Pass", comment="The only state")
state_machine = StateMachine(start_state=pass_state,
comment="My state machine",
version="1.1")
compiled_path = tmp_path / "state_machine.json"
state_machine.to_json(compiled_path)
with compiled_path.open() as fp:
compiled = json.load(fp)
assert compiled == {
"StartAt": pass_state.name,
"Comment": state_machine.comment,
"Version": state_machine.version,
"States": {
pass_state.name: {
"Comment": pass_state.comment,
"Type": "Pass",
"End": True,
},
},
}
def test_input_output_paths(capture_stdout):
input_path = "$.dataset2"
output_path = "$.val1"
pass_state = PassState("Pass 1", input_path=input_path, output_path=output_path)
state_machine = StateMachine(start_state=pass_state)
assert state_machine.compile() == {
"StartAt": "Pass 1",
"States": {
"Pass 1": {
"Type": "Pass",
"InputPath": "$.dataset2",
"OutputPath": "$.val1",
"End": True,
}
},
}
stdout = capture_stdout(
lambda: state_machine.simulate(
{
"comment": "Example for InputPath.",
"dataset1": {"val1": 1, "val2": 2, "val3": 3},
"dataset2": {"val1": "a", "val2": "b", "val3": "c"},
}
)
)
assert (
stdout
== """Starting simulation of state machine
Executing PassState('Pass 1')
State input: {'comment': 'Example for InputPath.', 'dataset1': {'val1': 1, 'val2': 2, 'val3': 3}, 'dataset2': {'val1': 'a', 'val2': 'b', 'val3': 'c'}}
State input after applying input path of $.dataset2: {'val1': 'a', 'val2': 'b', 'val3': 'c'}
Output from applying result path of $: {'val1': 'a', 'val2': 'b', 'val3': 'c'}
State output after applying output path of $.val1: a
State output: a
Terminating simulation of state machine
"""
)
def test_state_name_too_long():
with pytest.raises(AWSStepFuncsValueError,
match="State name cannot exceed 128 characters"):
PassState("a" * 129)
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
"""
)