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_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_fail_state(capture_stdout):
fail_state = FailState("Failure", error="IFailed", cause="I failed!")
state_machine = StateMachine(start_state=fail_state)
state_machine.compile() == {
"Type": "Fail",
"Error": "IFailed",
"Cause": "I failed!",
}
stdout = capture_stdout(lambda: state_machine.simulate())
assert (stdout == """Starting simulation of state machine
Executing FailState('Failure', error='IFailed', cause='I failed!')
State input: {}
FailStateError encountered in state
Checking for catchers
State output: {}
Terminating simulation of state machine
""")
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_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_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_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_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_map_state_foo(resource, iterator, state_input, capture_stdout,
mock_fn):
map_state = MapState(
"Validate-All",
input_path="$.detail",
items_path="$.shipped",
max_concurrency=0,
iterator=iterator,
)
state_machine = StateMachine(start_state=map_state)
stdout = capture_stdout(lambda: state_machine.simulate(
state_input,
resource_to_mock_fn={resource: mock_fn},
))
assert (stdout == """Starting simulation of state machine
Executing MapState('Validate-All')
State input: {'ship-date': '2016-03-14T01:59:00Z', 'detail': {'delivery-partner': 'UQS', 'shipped': [{'prod': 'R31', 'dest-code': 9511, 'quantity': 1344}, {'prod': 'S39', 'dest-code': 9511, 'quantity': 40}]}}
State input after applying input path of $.detail: {'delivery-partner': 'UQS', 'shipped': [{'prod': 'R31', 'dest-code': 9511, 'quantity': 1344}, {'prod': 'S39', 'dest-code': 9511, 'quantity': 40}]}
Items after applying items_path of $.shipped: [{'prod': 'R31', 'dest-code': 9511, 'quantity': 1344}, {'prod': 'S39', 'dest-code': 9511, 'quantity': 40}]
Starting simulation of state machine
Executing TaskState('Validate')
State input: {'prod': 'R31', 'dest-code': 9511, 'quantity': 1344}
State input after applying input path of $: {'prod': 'R31', 'dest-code': 9511, 'quantity': 1344}
Output from applying result path of $: {'prod': 'R31', 'dest-code': 9511, 'quantity': 2688}
State output after applying output path of $: {'prod': 'R31', 'dest-code': 9511, 'quantity': 2688}
State output: {'prod': 'R31', 'dest-code': 9511, 'quantity': 2688}
Terminating simulation of state machine
Starting simulation of state machine
Executing TaskState('Validate')
State input: {'prod': 'S39', 'dest-code': 9511, 'quantity': 40}
State input after applying input path of $: {'prod': 'S39', 'dest-code': 9511, 'quantity': 40}
Output from applying result path of $: {'prod': 'S39', 'dest-code': 9511, 'quantity': 80}
State output after applying output path of $: {'prod': 'S39', 'dest-code': 9511, 'quantity': 80}
State output: {'prod': 'S39', 'dest-code': 9511, 'quantity': 80}
Terminating simulation of state machine
Output from applying result path of $: [{'prod': 'R31', 'dest-code': 9511, 'quantity': 2688}, {'prod': 'S39', 'dest-code': 9511, 'quantity': 80}]
State output after applying output path of $: [{'prod': 'R31', 'dest-code': 9511, 'quantity': 2688}, {'prod': 'S39', 'dest-code': 9511, 'quantity': 80}]
State output: [{'prod': 'R31', 'dest-code': 9511, 'quantity': 2688}, {'prod': 'S39', 'dest-code': 9511, 'quantity': 80}]
Terminating simulation of state machine
""")
assert state_machine.compile() == {
"StartAt": "Validate-All",
"States": {
"Validate-All": {
"Type": "Map",
"InputPath": "$.detail",
"End": True,
"ItemsPath": "$.shipped",
"MaxConcurrency": 0,
"Iterator": {
"StartAt": "Validate",
"States": {
"Validate": {
"Type": "Task",
"End": True,
"Resource": "<arn>"
}
},
},
}
},
}