def test_add_waiting_for() -> None:
"""Test the ability to let a node wait for other nodes."""
run_graph = RunGraph()
run_graph.add_function_node("1")
run_graph.add_function_node("2")
run_graph.add_function_node("3")
run_graph.add_waiting_for("2", "1")
run_graph.add_waiting_for("3", ("2", "1"))
expected = ("1", "2", "3")
assert expected == run_graph.get_waiting_for_tuple("1")
def make_graph() -> RunGraph:
"""
Yield a simple graph.
Returns
-------
run_graph : RunGraph
A simple graph
"""
run_graph = RunGraph()
for i in range(6):
run_graph.add_function_node(str(i))
run_graph.add_waiting_for("4", "3")
run_graph.add_waiting_for("5", "3")
run_graph.add_waiting_for("3", "2")
run_graph.add_waiting_for("2", "0")
run_graph.add_waiting_for("1", "0")
return run_graph
def assert_waiting_for_graph(
node_zero_submitter: AbstractSubmitter,
node_one_submitter: AbstractSubmitter,
node_two_submitter: AbstractSubmitter,
node_three_submitter: AbstractSubmitter,
save_dir,
) -> None:
"""
Assert that the graph is running in correct order.
Parameters
----------
node_zero_submitter : AbstractSubmitter
Submitter object for node one
node_one_submitter : AbstractSubmitter
Submitter object for node one
node_two_submitter : AbstractSubmitter
Submitter object for node one
node_three_submitter : AbstractSubmitter
Submitter object for node one
save_dir : Path
Path to where the job artifacts are stored
"""
graph = RunGraph()
graph.add_function_node(
"node_zero",
{
"function": node_zero,
"args": (save_dir, ),
"kwargs": None
},
save_dir.joinpath("node_zero.py"),
node_zero_submitter,
)
graph.add_function_node(
"node_one",
{
"function": node_one,
"args": (save_dir, ),
"kwargs": None
},
save_dir.joinpath("node_one.py"),
node_one_submitter,
)
graph.add_function_node(
"node_two",
{
"function": node_two,
"args": (save_dir, ),
"kwargs": None
},
save_dir.joinpath("node_two.py"),
node_two_submitter,
)
graph.add_function_node(
"node_three",
{
"function": node_three,
"args": (save_dir, ),
"kwargs": None
},
save_dir.joinpath("node_three.py"),
node_three_submitter,
)
graph.add_waiting_for("node_two", "node_one")
graph.add_waiting_for("node_three", ("node_one", "node_two"))
runner = BoutRunner(graph)
runner.run()
node_three_submitter.wait_until_completed()
assert save_dir.joinpath("node_three.log").is_file()