def test_clear_operation_links_end(self, graph, sum_op, square_op,
negative_op):
graph.add_operations(sum_op, square_op, negative_op)
graph.add_link(sum_op, square_op, "sum", "n")
graph.add_link(square_op, negative_op, "square", "num")
graph.clear_operation_links(negative_op)
assert graph.links() == [(sum_op, square_op, "sum", "n")]
def test_remove_operation_linked(self, graph, sum_op, square_op, negative_op):
graph.add_operations(sum_op, square_op, negative_op)
graph.add_link(sum_op, square_op, "sum", "n")
graph.add_link(square_op, negative_op, "square", "num")
graph.remove_operation(sum_op)
assert graph.operations == [square_op, negative_op]
assert graph.links() == [(square_op, negative_op, "square", "num")]
def test_remove_link(self, graph, sum_op, square_op):
graph.add_operations(sum_op, square_op)
graph.add_link(sum_op, square_op, source_param="sum", dest_param="n")
graph.remove_link(sum_op, square_op, "sum", "n")
assert len(graph.get_inbound_links(sum_op)) == 0 and len(
graph.get_inbound_links(square_op)) == 0
assert len(graph.get_outbound_links(sum_op)) == 0 and len(
graph.get_outbound_links(square_op)) == 0
def test_remove_false(self, graph, a_op, b_op, c_op):
graph.add_operations(a_op, b_op, c_op)
link1 = (a_op, b_op, "n", "n")
link2 = (b_op, c_op, "n", "n")
graph.add_link(*link1)
graph.add_link(*link2)
return_value = graph.set_disabled(a_op, remove_orphan_links=False)
assert graph.disabled(a_op) is True
assert graph.links() == [link2]
assert return_value == [link1]
def test_value_false(self, graph, a_op, b_op, c_op):
graph.add_operations(a_op, b_op, c_op)
link1 = (a_op, b_op, "n", "n")
link2 = (b_op, c_op, "n", "n")
graph.add_link(*link1)
graph.add_link(*link2)
return_value = graph.set_disabled(sum_op, value=False)
assert graph.disabled(sum_op) is False
assert graph.links() == [link1, link2]
assert return_value == []
def test_value_and_remove_false_unmatched_names(self, graph, sum_op, square_op, negative_op):
graph.add_operations(sum_op, square_op, negative_op)
link1 = (sum_op, square_op, "sum", "n")
link2 = (square_op, negative_op, "square", "num")
graph.add_link(*link1)
graph.add_link(*link2)
return_value = graph.set_disabled(sum_op, value=False, remove_orphan_links=False)
assert graph.disabled(sum_op) is False
assert graph.links() == []
assert return_value == []
def test_default_unmatched_names(self, graph, sum_op, square_op, negative_op):
graph.add_operations(sum_op, square_op, negative_op)
link1 = (sum_op, square_op, "sum", "n")
link2 = (square_op, negative_op, "square", "num")
graph.add_link(*link1)
graph.add_link(*link2)
return_value = graph.set_disabled(sum_op)
assert graph.disabled(sum_op) is True
assert graph.links() == []
assert return_value == []
def test_set_disabled_value_false(self, graph, sum_op, square_op,
negative_op):
graph.add_operations(sum_op, square_op, negative_op)
link1 = (sum_op, square_op, "sum", "n")
link2 = (square_op, negative_op, "square", "num")
graph.add_link(*link1)
graph.add_link(*link2)
return_value = graph.set_disabled(sum_op, value=False)
assert graph.disabled(sum_op) is False
assert graph.links() == [link1, link2]
assert return_value == []
def test_links_multiple(self, graph, sum_op):
def my_func(x: int, y: int) -> (int, int):
return y, x
my_op = operation(my_func, output_names=("y", "x"))()
graph.add_operations(sum_op, my_op)
link1 = (my_op, sum_op, "y", "n1")
link2 = (my_op, sum_op, "x", "n2")
graph.add_link(*link1)
graph.add_link(*link2)
assert graph.links() == [link1, link2]
def test_default_with_links(self, graph, a_op, b_op, c_op):
graph = Graph()
graph.add_operations(a_op, b_op, c_op)
link_ab = (a_op, b_op, "n", "n")
link_bc = (b_op, c_op, "n", "n")
graph.add_link(*link_ab)
graph.add_link(*link_bc)
orphan_links = graph.set_disabled(a_op)
assert graph.disabled(a_op) is True
assert graph.disabled(b_op) is False
assert graph.disabled(c_op) is False
assert graph.links() == [link_bc]
assert orphan_links == []
def test_add_link(self, graph, sum_op, square_op):
graph.add_operation(sum_op)
graph.add_operation(square_op)
graph.add_link(source=sum_op, dest=square_op, source_param="sum", dest_param="n")
sum_inbound_links = graph.get_inbound_links(sum_op)
sum_outbound_links = graph.get_outbound_links(sum_op)
square_inbound_links = graph.get_inbound_links(square_op)
square_outbound_links = graph.get_outbound_links(square_op)
links = [("sum", "n")]
assert sum_inbound_links[square_op] == []
assert sum_outbound_links[square_op] == links
assert square_inbound_links[sum_op] == links
assert square_outbound_links[sum_op] == []
def test_toggle_with_links(self, graph, a_op, b_op, c_op):
graph.add_operations(a_op, b_op, c_op)
link1 = (a_op, b_op, "n", "n")
link2 = (b_op, c_op, "n", "n")
graph.add_link(*link1)
graph.add_link(*link2)
return_value = graph.toggle_disabled(a_op)
assert graph.disabled(a_op) is True
assert return_value == []
assert graph.links() == [link2]
return_value = graph.toggle_disabled(a_op)
assert graph.disabled(a_op) is False
assert return_value == []
assert graph.links() == [link1, link2]
def test_toggle_disabled_remove_false(self, graph, sum_op, square_op,
negative_op):
graph.add_operations(sum_op, square_op, negative_op)
link1 = (sum_op, square_op, "sum", "n")
link2 = (square_op, negative_op, "square", "num")
graph.add_link(*link1)
graph.add_link(*link2)
return_value = graph.toggle_disabled(sum_op, remove_orphan_links=False)
assert graph.disabled(sum_op) is True
assert return_value == graph.operation_links(sum_op)
assert graph.links() == [link1, link2]
return_value = graph.toggle_disabled(sum_op, remove_orphan_links=False)
assert graph.disabled(sum_op) is False
assert return_value == []
assert graph.links() == [link1, link2]
def test_toggle_remove_false_unmatched_names(self, graph, sum_op, square_op, negative_op):
graph.add_operations(sum_op, square_op, negative_op)
link1 = (sum_op, square_op, "sum", "n")
link2 = (square_op, negative_op, "square", "num")
graph.add_link(*link1)
graph.add_link(*link2)
graph._pretty_print()
return_value = graph.toggle_disabled(sum_op, remove_orphan_links=False)
graph._pretty_print()
assert graph.disabled(sum_op) is True
assert return_value == [link1]
assert graph.links() == []
return_value = graph.toggle_disabled(sum_op, remove_orphan_links=False)
assert graph.disabled(sum_op) is False
assert return_value == []
assert graph.links() == []
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_add_multiple_links(self, graph, sum_op, square_op):
graph.add_operation(sum_op)
graph.add_operation(square_op)
import math
def my_sqrt(num):
return math.sqrt(num)
sqrt_op = operation(my_sqrt, output_names=("sqrt"))()
graph.add_operation(sqrt_op)
# sum -> square -> sqrt
# \ |
# -------->------
graph.add_link(source=sum_op, dest=square_op, source_param="sum", dest_param="n")
graph.add_link(source=square_op, dest=sqrt_op, source_param="square", dest_param="num")
graph.add_link(source=sum_op, dest=sqrt_op, source_param="sum", dest_param="num")
# sum -> square
assert graph.get_inbound_links(sum_op)[square_op] == []
assert graph.get_outbound_links(sum_op)[square_op] == [("sum", "n")]
assert graph.get_inbound_links(square_op)[sum_op] == [("sum", "n")]
assert graph.get_outbound_links(square_op)[sum_op] == []
# square -> sqrt
assert graph.get_inbound_links(square_op)[sqrt_op] == []
assert graph.get_outbound_links(square_op)[sqrt_op] == [("square", "num")]
assert graph.get_inbound_links(sqrt_op)[square_op] == [("square", "num")]
assert graph.get_outbound_links(sqrt_op)[square_op] == []
# sum -> sqrt
assert graph.get_inbound_links(sum_op)[sqrt_op] == []
assert graph.get_outbound_links(sum_op)[sqrt_op] == [("sum", "num")]
assert graph.get_inbound_links(sqrt_op)[sum_op] == [("sum", "num")]
assert graph.get_outbound_links(sqrt_op)[sum_op] == []
def test_operation_links_multiple(self, graph, sum_op, square_op,
negative_op):
def my_func(x: int, y: int) -> (int, int):
return y, x
my_op = operation(my_func, output_names=("y", "x"))()
graph.add_operations(sum_op, square_op, negative_op, my_op)
link1 = (my_op, sum_op, "y", "n1")
link2 = (my_op, sum_op, "x", "n2")
link3 = (sum_op, square_op, "sum", "n")
link4 = (square_op, negative_op, "square", "num")
graph.add_link(*link1)
graph.add_link(*link2)
graph.add_link(*link3)
graph.add_link(*link4)
assert graph.operation_links(my_op) == [link1, link2]
assert graph.operation_links(sum_op) == [link3]
assert graph.operation_links(negative_op) == []
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_add_link_bad_dest(self, graph, sum_op):
with pytest.raises(TypeError):
graph.add_link(sum_op, "bad", "", "")
def test_add_link_messing_dest(self, graph, sum_op, square_op):
graph.add_operation(sum_op)
with pytest.raises(ValueError):
graph.add_link(sum_op, square_op, "sum", "n")
def test_add_link_missing_dest_param(self, graph, sum_op, square_op):
with pytest.raises(ValueError):
graph.add_link(source=sum_op,
dest=square_op,
source_param="dne",
dest_param="n")
def test_get_outbound_links(self, graph, sum_op, square_op):
graph.add_operations(sum_op, square_op)
graph.add_link(sum_op, square_op, "sum", "n")
assert graph.get_outbound_links(square_op) == {}
assert graph.get_outbound_links(sum_op) == {square_op: [("sum", "n")]}
def test_operation_links(self, graph, sum_op, square_op):
graph.add_operations(sum_op, square_op)
link = (sum_op, square_op, "sum", "n")
graph.add_link(*link)
assert graph.operation_links(sum_op) == [link]
assert graph.operation_links(square_op) == []