DeduceNodeLocationsFunction implements Function<MutableNetwork<Node, Edge>, MutableNetwork<Node, Edge>> { @Override public MutableNetwork<Node, Edge> apply(MutableNetwork<Node, Edge> network) { Networks.replaceDirectedNetworkNodes( network, (Node node) -> { if (!isDeducible(node)) { return node; } ParallelInstructionNode castNode = ((ParallelInstructionNode) node); ExecutionLocation location; if (executesInSdkHarness(castNode.getParallelInstruction())) { location = ExecutionLocation.SDK_HARNESS; } else { location = ExecutionLocation.RUNNER_HARNESS; } return ParallelInstructionNode.create(castNode.getParallelInstruction(), location); }); return network; }  @Override MutableNetwork<Node, Edge> apply(MutableNetwork<Node, Edge> network);  }
@Test public void testEmptyNetwork() { assertEquals( createEmptyNetwork(), new DeduceNodeLocationsFunction().apply(createEmptyNetwork())); }
@Test public void testSingleNodeWithSdkRead() throws Exception { Node unknown = createReadNode("Unknown", CUSTOM_SOURCE); MutableNetwork<Node, Edge> network = createEmptyNetwork(); network.addNode(unknown); Network<Node, Edge> inputNetwork = ImmutableNetwork.copyOf(network); network = new DeduceNodeLocationsFunction().apply(network); assertThatNetworksAreIdentical(inputNetwork, network); for (Node node : ImmutableList.copyOf(network.nodes())) { assertNodesIdenticalExceptForExecutionLocation(unknown, node); assertThatLocationIsProperlyDeduced(node, ExecutionLocation.SDK_HARNESS); } }
@Test public void testSingleNodeWithRunnerRead() throws Exception { Node unknown = createReadNode("Unknown", RUNNER_SOURCE); MutableNetwork<Node, Edge> network = createEmptyNetwork(); network.addNode(unknown); Network<Node, Edge> inputNetwork = ImmutableNetwork.copyOf(network); network = new DeduceNodeLocationsFunction().apply(network); assertThatNetworksAreIdentical(inputNetwork, network); for (Node node : ImmutableList.copyOf(network.nodes())) { assertNodesIdenticalExceptForExecutionLocation(unknown, node); assertThatLocationIsProperlyDeduced(node, ExecutionLocation.RUNNER_HARNESS); } }
@Test public void testSingleNodeWithSdkParDo() throws Exception { Node unknown = createParDoNode("Unknown", DO_FN); MutableNetwork<Node, Edge> network = createEmptyNetwork(); network.addNode(unknown); Network<Node, Edge> inputNetwork = ImmutableNetwork.copyOf(network); network = new DeduceNodeLocationsFunction().apply(network); assertThatNetworksAreIdentical(inputNetwork, network); for (Node node : ImmutableList.copyOf(network.nodes())) { assertNodesIdenticalExceptForExecutionLocation(unknown, node); assertThatLocationIsProperlyDeduced(node, ExecutionLocation.SDK_HARNESS); } }
@Test public void testSingleNodeWithRunnerParDo() throws Exception { Node unknown = createParDoNode("Unknown", "RunnerDoFn"); MutableNetwork<Node, Edge> network = createEmptyNetwork(); network.addNode(unknown); Network<Node, Edge> inputNetwork = ImmutableNetwork.copyOf(network); network = new DeduceNodeLocationsFunction().apply(network); assertThatNetworksAreIdentical(inputNetwork, network); for (Node node : ImmutableList.copyOf(network.nodes())) { assertNodesIdenticalExceptForExecutionLocation(unknown, node); assertThatLocationIsProperlyDeduced(node, ExecutionLocation.RUNNER_HARNESS); } }
@Test public void testMultipleNodesDeduced() throws Exception { Node a = createReadNode("A", CUSTOM_SOURCE); Node b = createReadNode("B", RUNNER_SOURCE); Node c = createParDoNode("C", "RunnerDoFn"); Node d = createParDoNode("D", DO_FN); Node e = createParDoNode("E", DO_FN); MutableNetwork<Node, Edge> network = createEmptyNetwork(); network.addNode(a); network.addNode(b); network.addNode(c); network.addNode(d); network.addNode(e); network.addEdge(a, e, DefaultEdge.create()); network.addEdge(b, e, DefaultEdge.create()); network.addEdge(e, c, DefaultEdge.create()); network.addEdge(e, d, DefaultEdge.create()); Network<Node, Edge> inputNetwork = ImmutableNetwork.copyOf(network); network = new DeduceNodeLocationsFunction().apply(network); assertThatNetworksAreIdentical(inputNetwork, network); assertAllNodesDeducedExceptFlattens(network); }
@Test public void testGraphWithNonDeducibleNodes() throws Exception { Node a = createReadNode("A", CUSTOM_SOURCE); Node out1 = InstructionOutputNode.create(new InstructionOutput(), "fakeId"); Node b = createReadNode("B", RUNNER_SOURCE); Node out2 = InstructionOutputNode.create(new InstructionOutput(), "fakeId"); Node c = createParDoNode("C", "RunnerDoFn"); Node flatten = ParallelInstructionNode.create( new ParallelInstruction().setName("Flatten").setFlatten(new FlattenInstruction()), Nodes.ExecutionLocation.UNKNOWN); Node d = createParDoNode("D", DO_FN); MutableNetwork<Node, Edge> network = createEmptyNetwork(); network.addNode(a); network.addNode(out1); network.addNode(b); network.addNode(out2); network.addNode(c); network.addNode(flatten); network.addNode(d); network.addEdge(a, out1, DefaultEdge.create()); network.addEdge(b, out2, DefaultEdge.create()); network.addEdge(out1, flatten, DefaultEdge.create()); network.addEdge(out2, flatten, DefaultEdge.create()); network.addEdge(out2, c, DefaultEdge.create()); network.addEdge(flatten, d, DefaultEdge.create()); Network<Node, Edge> inputNetwork = ImmutableNetwork.copyOf(network); network = new DeduceNodeLocationsFunction().apply(network); assertThatNetworksAreIdentical(inputNetwork, network); assertAllNodesDeducedExceptFlattens(network); }