def test_dependency_set_computations(): """ Test methods for computing upstream and downstream dependency sets for specific nodes in the graph """ 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]) assert _GraphUtil.downstream_dependency_set(node0, graph) == frozenset( [node2, node3]) assert _GraphUtil.downstream_dependency_set(node1, graph) == frozenset( [node2, node3]) assert _GraphUtil.upstream_dependency_set(node2, graph) == frozenset( [node0, node1]) assert _GraphUtil.upstream_dependency_set(node3, graph) == frozenset( [node0, node1])
def test_bundler_inserter(fc_node_builder): """ Make sure the insert_bundler_nodes() method works properly for a general graph, with multiple upstream dependency sets """ 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': ['test1', 'test2'], }) node4 = OperatorNode({}, { 'name': 'test4', 'type': 'none', 'upstream_dependencies': ['test2'], }) graph = _GraphUtil.build_subgraph([node0, node1, node2, node3, node4]) dep_set_0 = frozenset([node1, node2]) dep_set_1 = frozenset([node0]) bundlers = _GraphUtil.insert_bundler_nodes([node2, node3], graph, fc_node_builder=fc_node_builder) assert set(bundlers.keys()) == set([dep_set_0, dep_set_1]) bundler_0 = bundlers[dep_set_0] assert bundler_0.name == 'test1-test2-322ef2-fc' assert _GraphUtil.upstream_dependency_set(bundler_0, graph) == dep_set_0 assert _GraphUtil.downstream_dependency_set(bundler_0, graph) == frozenset([node3]) bundler_1 = bundlers[dep_set_1] assert bundler_1.name == 'test0-f6f406-fc' assert _GraphUtil.upstream_dependency_set(bundler_1, graph) == dep_set_1 assert _GraphUtil.downstream_dependency_set(bundler_1, graph) == frozenset([node2])
def test_prune_nodes(): """ Test to make sure that node-pruning works """ node0 = OperatorNode({}, {'name': 'test0', 'type': 'none'}) node1 = OperatorNode({}, { 'name': 'test1', 'type': 'none', 'upstream_dependencies': ['test0'] }) 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.prune_nodes(graph, node_selector=lambda node: node.name == 'test1') assert _GraphUtil.upstream_dependency_set(node3, graph) == frozenset([node0]) assert _GraphUtil.downstream_dependency_set(node0, graph) == frozenset( [node2, node3])
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])
def test_default_bundler_inserter(fc_node_builder): """ Make sure the insert_bundler_nodes() method works properly for a simpler graph, with a single distinct upstream dependency set and a specified default bundler node """ 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': ['test1', 'test2'], }) resource = OperatorNode({}, { 'name': 'cluster', 'type': 'none', }) graph = _GraphUtil.build_subgraph([node0, node1, node2, node3]) dep_set = frozenset([node1, node2]) bundlers = _GraphUtil.insert_bundler_nodes([node3], graph, default_bundler_node=resource, fc_node_builder=fc_node_builder) assert set(bundlers.keys()) == set([dep_set]) bundler_0 = bundlers[dep_set] assert bundler_0 == resource assert _GraphUtil.upstream_dependency_set(bundler_0, graph) == dep_set assert _GraphUtil.downstream_dependency_set(bundler_0, graph) == frozenset([node3])
def test_attach_destroy_resource_no_sentinel_because_other_leaf_nodes( fc_node_builder): """ Make sure downstream resource-delete operator insertion works correctly """ 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': ['test2'], 'requires_resources': ['cluster'], }) node4 = OperatorNode({}, { 'name': 'test4', 'type': 'none', 'upstream_dependencies': ['test2', 'test1'], 'requires_resources': ['cluster'], }) node5 = OperatorNode({}, { 'name': 'test4', 'type': 'none', 'upstream_dependencies': ['test3'], }) graph = _GraphUtil.build_subgraph( [node0, node1, node2, node3, node4, node5]) destroy_resource_node = OperatorNode({}, { 'name': 'destroy_cluster', 'type': 'none', }) resource = MockResource(name='cluster', create_operator=None, destroy_operator=destroy_resource_node, disable_sentinel_node=False) _GraphUtil.attach_destroy_resource(resource, graph, fc_node_builder) destroy_node_upstream = _GraphUtil.upstream_dependency_set( destroy_resource_node, graph) assert destroy_node_upstream == frozenset([node3, node4]) assert len(_GraphUtil.downstream_dependency_set(node4, graph)) == 2 assert destroy_resource_node in _GraphUtil.downstream_dependency_set( node4, graph) # sentinel node should also have been inserted assert any(item.name == 'destroy_cluster-sentinel' for item in _GraphUtil.downstream_dependency_set(node4, graph))
def test_attach_destroy_resource_no_sentinel_by_config(fc_node_builder): """ Make sure downstream resource-delete operator insertion works correctly """ 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': ['test2'], 'requires_resources': ['cluster'], }) node4 = OperatorNode({}, { 'name': 'test4', 'type': 'none', 'upstream_dependencies': ['test2', 'test1'], 'requires_resources': ['cluster'], }) node5 = OperatorNode({}, { 'name': 'test4', 'type': 'none', 'upstream_dependencies': ['test3'], 'requires_resources': ['cluster'], }) graph = _GraphUtil.build_subgraph( [node0, node1, node2, node3, node4, node5]) destroy_resource_node = OperatorNode({}, { 'name': 'destroy_cluster', 'type': 'none', }) resource = MockResource(name='cluster', create_operator=None, destroy_operator=destroy_resource_node, disable_sentinel_node=True) _GraphUtil.attach_destroy_resource(resource, graph, fc_node_builder) destroy_node_upstream = _GraphUtil.upstream_dependency_set( destroy_resource_node, graph) assert destroy_node_upstream == frozenset([node4, node5]) assert _GraphUtil.downstream_dependency_set(node4, graph) == \ _GraphUtil.downstream_dependency_set(node5, graph) == \ frozenset([destroy_resource_node])
def test_attach_create_resource(fc_node_builder): """ Make sure upstream resource-create operator insertion works correctly """ 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': ['test2'], 'requires_resources': ['cluster'], }) node4 = OperatorNode({}, { 'name': 'test4', 'type': 'none', 'upstream_dependencies': ['test2', 'test1'], 'requires_resources': ['cluster'], }) node5 = OperatorNode({}, { 'name': 'test4', 'type': 'none', 'upstream_dependencies': ['test2'], 'requires_resources': ['cluster'], }) graph = _GraphUtil.build_subgraph( [node0, node1, node2, node3, node4, node5]) create_resource_node = OperatorNode({}, { 'name': 'create_cluster', 'type': 'none', }) resource = MockResource(name='cluster', create_operator=create_resource_node, destroy_operator=None, disable_sentinel_node=None) _GraphUtil.attach_create_resource(resource, graph, fc_node_builder) assert set(_GraphUtil.upstream_resource_boundary('cluster', graph)) == set( [node3, node4, node5]) cluster_upstream_deps = _GraphUtil.upstream_dependency_set( create_resource_node, graph) # There should be 2 bundlers because there are 2 distinct dependency sets among # the nodes in the upstream resource boundary assert len(cluster_upstream_deps) == 2 node3_upstream = _GraphUtil.upstream_dependency_set(node3, graph) assert len(node3_upstream) == 2 assert _GraphUtil.upstream_dependency_set(node5, graph) == node3_upstream assert create_resource_node in node3_upstream node3_bundler = list(node3_upstream - frozenset([create_resource_node]))[0] assert _GraphUtil.upstream_dependency_set(node3_bundler, graph) == frozenset([node2]) node4_upstream = _GraphUtil.upstream_dependency_set(node4, graph) assert len(node4_upstream) == 2 assert create_resource_node in node4_upstream node4_bundler = list(node4_upstream - frozenset([create_resource_node]))[0] assert _GraphUtil.upstream_dependency_set( node4_bundler, graph) == frozenset([node1, node2]) create_resource_downstream_deps = _GraphUtil.downstream_dependency_set( create_resource_node, graph) assert create_resource_downstream_deps == frozenset([node3, node4, node5])