def test_flow_control_inserter_invalid_args(fc_node_builder):
""" Make sure that attach_flow_control_between() raises an exception if
there are any missing dependency linkages between the upstream and
downstream nodes
"""
node0 = OperatorNode({}, {'name': 'test0', 'type': 'none'})
node1 = OperatorNode({}, {'name': 'test1', 'type': 'none'})
node2 = OperatorNode({}, {
'name': 'test2',
'type': 'none',
'upstream_dependencies': ['test0']
})
node3 = OperatorNode({}, {
'name': 'test3',
'type': 'none',
'upstream_dependencies': ['test0', 'test1']
})
graph = _GraphUtil.build_subgraph([node0, node1, node2, node3])
with pytest.raises(InvalidFlowControlNode):
_GraphUtil.attach_flow_control_between(
upstream=frozenset([node0, node1]),
downstream=frozenset([node2, node3]),
graph=graph,
fc_node_builder=fc_node_builder,
)
def test_flow_control_inserter_auto_fc_node(fc_node_builder):
""" Test method for inserting flow control nodes between upstream and downstream
nodes with auto-created flow-control node
"""
node0 = OperatorNode({}, {'name': 'test0', 'type': 'none'})
node1 = OperatorNode({}, {'name': 'test1', 'type': 'none'})
node2 = OperatorNode({}, {
'name': 'test2',
'type': 'none',
'upstream_dependencies': ['test0', 'test1']
})
node3 = OperatorNode({}, {
'name': 'test3',
'type': 'none',
'upstream_dependencies': ['test0', 'test1']
})
graph = _GraphUtil.build_subgraph([node0, node1, node2, node3])
_GraphUtil.attach_flow_control_between(
upstream=frozenset([node0, node1]),
downstream=frozenset([node2, node3]),
graph=graph,
fc_node_builder=fc_node_builder,
)
assert len(graph.nodes()) == 5
fc_node = [node for node in graph.nodes() if node.name.endswith('-fc')]
assert len(fc_node) == 1
assert fc_node[0].type == NodeTypes.OPERATOR
assert fc_node[0].name == "test2-test3-d454d5-fc"
assert _GraphUtil.upstream_dependency_set(fc_node[0], graph) == frozenset(
[node0, node1])
assert _GraphUtil.downstream_dependency_set(
fc_node[0], graph) == frozenset([node2, node3])
assert _GraphUtil.downstream_dependency_set(node0, graph) == frozenset(
[fc_node[0]])
assert _GraphUtil.downstream_dependency_set(node1, graph) == frozenset(
[fc_node[0]])
assert _GraphUtil.upstream_dependency_set(node2,
graph) == frozenset([fc_node[0]])
assert _GraphUtil.upstream_dependency_set(node3,
graph) == frozenset([fc_node[0]])
def test_flow_control_inserter():
""" Test method for inserting flow control nodes between upstream and downstream
nodes with preexisting flow control node
"""
node0 = OperatorNode({}, {'name': 'test0', 'type': 'none'})
node1 = OperatorNode({}, {'name': 'test1', 'type': 'none'})
node2 = OperatorNode({}, {
'name': 'test2',
'type': 'none',
'upstream_dependencies': ['test0', 'test1']
})
node3 = OperatorNode({}, {
'name': 'test3',
'type': 'none',
'upstream_dependencies': ['test0', 'test1']
})
fc_node = OperatorNode({}, {'name': 'my-fc-node', 'type': 'none'})
graph = _GraphUtil.build_subgraph([node0, node1, node2, node3])
_GraphUtil.attach_flow_control_between(
upstream=frozenset([node0, node1]),
downstream=frozenset([node2, node3]),
graph=graph,
fc_node=fc_node,
)
assert len(graph.nodes()) == 5
assert fc_node in graph.nodes()
assert _GraphUtil.upstream_dependency_set(fc_node, graph) == frozenset(
[node0, node1])
assert _GraphUtil.downstream_dependency_set(fc_node, graph) == frozenset(
[node2, node3])
assert _GraphUtil.downstream_dependency_set(node0,
graph) == frozenset([fc_node])
assert _GraphUtil.downstream_dependency_set(node1,
graph) == frozenset([fc_node])
assert _GraphUtil.upstream_dependency_set(node2,
graph) == frozenset([fc_node])
assert _GraphUtil.upstream_dependency_set(node3,
graph) == frozenset([fc_node])