def test_machine_accepts_fields_as_kwargs(example):
m1 = Machine.parse([],
comment="This is a comment",
version="1.0",
timeout_seconds=60)
assert m1.comment == "This is a comment"
assert m1.version == "1.0"
assert m1.timeout_seconds == 60
m2 = Machine.parse(example("hello_world"), timeout_seconds=35)
assert m2.timeout_seconds == 35
def test_inserts_and_removes_in_a_sequence():
sequence = Machine.parse(["a", "c", "e"])
assert sequence.dry_run() == ["a", "c", "e"]
sequence.insert("b", before="c")
assert sequence.dry_run() == ["a", "b", "c", "e"]
sequence.insert("f", after="e")
assert sequence.dry_run() == ["a", "b", "c", "e", "f"]
sequence.remove("c")
assert sequence.dry_run() == ["a", "b", "e", "f"]
sequence.remove("a")
assert sequence.dry_run() == ["b", "e", "f"]
sequence.remove("f")
assert sequence.dry_run() == ["b", "e"]
sequence.insert("x", before="b")
assert sequence.dry_run() == ["x", "b", "e"]
sequence.remove("b")
sequence.remove("e")
sequence.remove("x")
assert sequence.dry_run() == []
assert sequence.start_at is None
def test_create_state_machine(sfn_client):
sm = Machine.parse([
TActivities.a,
TActivities.b,
[
[TActivities.d, TActivities.e],
[TActivities.f],
],
[
[TActivities.g],
[TActivities.h],
],
])
activities: Dict[str, ActivityListItem] = {
a.name: a
for a in sfn_client.list_activities().activities
}
def state_visitor(state: State, compiled_state: Dict):
if compiled_state.get(
"Type") == "Task" and "Resource" not in compiled_state:
compiled_state["Resource"] = activities[state.name].activity_arn
sm_definition = sm.to_json(state_visitor=state_visitor)
sfn_client.create_state_machine(
name=f"{tests_profile.names_prefix}-something-statemachine",
definition=sm_definition,
role_arn=tests_profile.role_arn,
)
def test_simple_parallel():
source = [["a"], ["b"]]
s = Machine.parse(source)
assert len(s.states) == 1
assert isinstance(s.states[s.start_at], Parallel)
assert s.dry_run() == source
c = s.compile()
assert c["States"][c["StartAt"]]["Type"] == "Parallel"
def test_runs_hello_world_machine(example):
sm = Machine.parse(example("hello_world"))
runner = Runner(resources=ResourceManager(
providers={
"arn:aws:lambda:us-east-1:123456789012:function:HelloWorld":
lambda x: "Hello, world!"
}))
assert runner.run(sm) == (sm.start_at_state, "Hello, world!")
def test_machine_state(example):
source = example("hello_world")
machine = Machine.parse(source)
assert machine.type == States.Machine
assert isinstance(machine.start_at_state, Task)
assert machine.start_at == "Hello World"
assert machine.start_at_state.name == "Hello World"
assert machine.compile() == source
def test_executes_hello_world_state(example):
hello_world_state = Machine.parse(example("hello_world")).start_at_state
assert isinstance(hello_world_state, State)
resources = ResourceManager(
providers={
"arn:aws:lambda:us-east-1:123456789012:function:HelloWorld":
lambda x: "Hello, world!"
})
next_state, output = hello_world_state.execute({},
resource_resolver=resources)
assert output == "Hello, world!"
def test_visitor_visits_every_compiled_state_dictionary():
m = Machine.parse([
["a", "b", "c"],
["1", "2"],
])
def state_visitor(state: State, compiled_state: Dict):
compiled_state["Resource"] = f"arn.funny.{state.name}"
c = m.compile(state_visitor=state_visitor)
assert c["States"][c["StartAt"]]["Branches"][0]["States"]["a"][
"Resource"] == "arn.funny.a"
def test_inserts_and_removes_in_parallels():
sm = Machine.parse([
["b", "c"],
["1", "2"],
])
sm.start_at_state.branches[0].insert("a", before="b")
sm.start_at_state.branches[1].insert("3", after="2")
assert sm.dry_run() == [
["a", "b", "c"],
["1", "2", "3"],
]
def test_parallel_inside_sequence():
source = [
"a",
[
["b11", "b12"],
["b21", "b22"],
],
"c",
]
s = Machine.parse(source)
assert len(s.states) == 3
assert s.start_at == "a"
assert s.dry_run() == source
c = s.compile()
assert c["States"][c["States"]["a"]["Next"]]["Type"] == "Parallel"
def test_custom_task():
@dataclasses.dataclass
class Task:
name: str
def __str__(self):
return self.name
def get_state_attrs(self, state):
return {"Resource": f"arn:{self.name}"}
m = Machine.parse([
Task("first"),
Task("second"),
])
c = m.compile()
assert c["States"]["first"]["Type"] == "Task"
assert c["States"]["first"]["Resource"] == "arn:first"
assert c["States"]["second"]["Resource"] == "arn:second"
def test_parallel_inside_parallel():
source = [[
"a",
], [
[
[
"b11",
],
[
"b21",
],
],
"b3",
]]
s = Machine.parse(source)
assert s.dry_run() == source
c = s.compile()
assert c["States"][c["StartAt"]]["Type"] == "Parallel"
def test_runs_job_status_poller(example):
source = example("job_status_poller")
sm = Machine.parse(source)
assert sm.compile() == source
runner = Runner()
@runner.resource_provider("arn:aws:lambda:REGION:ACCOUNT_ID:function:SubmitJob")
def submit_job(payload):
return payload
@runner.resource_provider("arn:aws:lambda:REGION:ACCOUNT_ID:function:CheckJob")
def check_job(payload):
if random.random() < 0.2:
return "FAILED"
else:
return "SUCCEEDED"
final_state, output = runner.run(sm)
assert final_state.name in ("Get Final Job Status", "Job Failed")
def test_dicts_as_states():
m = Machine.parse([
{
"Comment": "first",
"Type": "Task",
"Resource": "arn:first",
},
{
"Name": "second",
"Type": "Pass",
"Resource": "arn:second",
},
"third",
])
c = m.compile()
assert c["StartAt"] == "arn:first"
assert len(c["States"]) == 3
assert c["States"]["arn:first"]["Comment"] == "first"
assert c["States"]["arn:first"]["Resource"] == "arn:first"
assert "Name" not in c["States"]["second"]
assert c["States"]["second"]["Type"] == "Pass"
def test_input_passed_to_next_task():
sm = Machine.parse([
{
"InputPath": "$.first_input",
"ResultPath": "$.first_output",
"Resource": "MultiplierByTwo",
},
{
"InputPath": "$.first_output",
"ResultPath": "$.second_output",
"Resource": "MultiplierByThree",
},
{
"Resource": "Validator",
},
])
runner = Runner()
runner.resource_provider("MultiplierByTwo")(lambda x: x * 2)
runner.resource_provider("MultiplierByThree")(lambda x: x * 3)
@runner.resource_provider("Validator")
def validate_input(input):
assert input == {
"first_input": 1111,
"first_output": 2222,
"second_output": 6666,
}
# NB!
return input
final_state, output = runner.run(sm, input={"first_input": 1111})
assert output == {
"first_input": 1111,
"first_output": 2222,
"second_output": 6666,
}
def test_runs_job_status_poller(example):
source = example("job_status_poller")
sm = Machine.parse(source)
assert sm.compile() == source
assert sm.max_pages == 2
assert sm.states['Submit Job'].parallel_pages == True
runner = Runner()
@runner.resource_provider("arn:aws:lambda:REGION:ACCOUNT_ID:function:SubmitJob")
def submit_job(payload):
return payload
@runner.resource_provider("arn:aws:lambda:REGION:ACCOUNT_ID:function:CheckJob")
def check_job(payload):
if payload < 30:
return "FAILED"
else:
return "SUCCEEDED"
final_state, output = runner.run(sm, {'input': 25})
assert final_state.name in "Job Failed"
final_state, output = runner.run(sm, {'input': 40})
assert final_state.name in "Consolidator Output"
def test_simple_sequence():
s = Machine.parse(["a", "b"])
assert len(s.states) == 2
assert s.start_at == "a"
assert s.dry_run() == ["a", "b"]
from aws_sfn_builder import Machine
machine = Machine.parse([
"a",
[
["b-10", "b-11"],
["b-20"],
],
"c",
])
print(machine.to_json())
def test_empty_machine():
m = Machine.parse([])
assert m.start_at is None
assert not m.states
assert m.dry_run() == []
def test_append_to_a_sequence():
sm = Machine.parse([])
sm.append("a")
sm.append("b")
assert sm.dry_run() == ["a", "b"]