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_operation_not_remove_orphan_links(self, graph, sum_op,
square_op):
graph.add_operations(sum_op, square_op)
to_remove = graph.remove_operation(square_op,
remove_orphan_links=False)
assert graph.operations == [sum_op]
assert to_remove == (graph.get_inbound_links(square_op),
graph.get_outbound_links(square_op))
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_auto_connect_all_only_matching_types(self, graph):
def my_increment(n: int) -> int:
return n + 1
def my_decrement(m: int) -> int:
return m - 1
increment_op = operation(my_increment, output_names=("increment", ))()
decrement_op = operation(my_decrement, output_names=("decrement", ))()
graph.add_operations(increment_op, decrement_op)
graph.auto_connect_all()
assert graph.links() == []
def test_auto_connect_all_matching_names_and_types(self, graph):
def my_increment(n: int) -> int:
return n + 1
def my_decrement(increment: int) -> int:
return increment - 1
increment_op = operation(my_increment, output_names=("increment",))()
decrement_op = operation(my_decrement, output_names=("end_result",))()
graph.add_operations(increment_op, decrement_op)
graph.auto_connect_all()
assert graph.links() == [(increment_op, decrement_op, "increment", "increment")]
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_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_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_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_clear_operation_links_unlinked_operation(self, graph, sum_op,
square_op):
# TODO should this raise an exception?
graph.add_operations(sum_op, square_op)
graph.clear_operation_links(sum_op)
assert graph.links() == []
def test_remove_link_no_links(self, graph, sum_op, square_op):
graph.add_operations(sum_op, square_op)
with pytest.raises(ValueError):
graph.remove_link(sum_op, square_op, "sum", "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_remove_operation(self, graph, sum_op, square_op):
graph.add_operations(sum_op, square_op)
return_value = graph.remove_operation(sum_op)
assert graph.operations == [square_op]
assert return_value is None # no return by default
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) == []
def test_add_operations(self, graph, sum_op, square_op):
operations = [sum_op, square_op]
graph.add_operations(*operations)
assert graph.operations == operations
def test_auto_connect_all_none_matching(self, graph, sum_op, square_op,
negative_op):
graph.add_operations(sum_op, square_op, negative_op)
graph.auto_connect_all()
assert graph.links() == []
def test_clear_operations(self, graph, sum_op, square_op):
graph.add_operations(sum_op, square_op)
graph.clear_operations()
assert graph.operations == []