def test_as_dask_graph_multiple_links(self, graph, sum_op, square_op,
negative_op):
def my_func(x: int, y: int) -> (int, int):
return y, x
# Connect sum_op to square_op.
# Connect sum_op to my_op's x, square_op to my_op's y.
# Leave negative_op unconnected
my_op = operation(my_func, output_names=("y", "x"))()
graph.add_operations(sum_op, square_op, negative_op, my_op)
graph.add_link(sum_op, square_op, "sum", "n")
graph.add_link(sum_op, my_op, "sum", "x")
graph.add_link(square_op, my_op, "square", "y")
dask_graph, end_ids = graph.as_dask_graph()
# Should look like:
sum_wrapper = dask_graph["0"]
square_wrapper = dask_graph["1"]
negative_wrapper = dask_graph["2"]
my_wrapper = dask_graph["3"]
# sum_op has no dependent nodes (no ops connect into it)
assert len(sum_wrapper) == 1
assert sum_wrapper[0].node is sum_op
# square_op has 1 dependent node, takes sum_op's output
assert len(square_wrapper) == 2
assert square_wrapper[0].node is square_op
# negative_op has no dependent nodes; is unconnected
assert len(negative_wrapper) == 1
assert negative_wrapper[0].node is negative_op
# my_op has two dependent nodes; sum_op and square_op connect to its inputs
assert len(my_wrapper) == 3
assert my_wrapper[0].node is my_op
assert my_wrapper[1] == "0" # sum_op
assert my_wrapper[2] == "1" # square_op
# negative_op, and my_op should be end nodes
assert sorted(end_ids) == sorted(["2", "3"])
def test_as_dask_graph(self, graph, sum_op, square_op, negative_op):
# Connect sum_op to square_op; don't connect negative_op
graph.add_operations(sum_op, square_op, negative_op)
graph.add_link(sum_op, square_op, "sum", "n")
dask_graph, end_ids = graph.as_dask_graph()
# Should look like:
# { "0": (<sum_op>,), "1": (<square_op>, "0"), "2": (<negative_op>,) }
sum_wrapper = dask_graph["0"]
square_wrapper = dask_graph["1"]
negative_wrapper = dask_graph["2"]
assert len(sum_wrapper) == 1
assert sum_wrapper[0].node is sum_op
assert len(square_wrapper) == 2
assert square_wrapper[0].node is square_op
assert square_wrapper[1] == "0"
assert len(negative_wrapper) == 1
assert negative_wrapper[0].node is negative_op
# Both square_op and negative_op should be end nodes
assert sorted(end_ids) == sorted(["1", "2"])
def test_as_dask_graph_no_links(self, graph, sum_op):
graph.add_operation(sum_op)
dask_graph, end_ids = graph.as_dask_graph()
assert len(dask_graph["0"]) == 1
assert dask_graph["0"][0].node is sum_op
assert end_ids == ["0"]
def test_as_dask_graph_empty(self, graph):
# Empty graph, no end nodes
assert graph.as_dask_graph() == ({}, [])