use of com.google.api.services.dataflow.model.InstructionOutput in project beam by apache.
the class RemoveFlattenInstructionsFunctionTest method testRemoveFlatten.
@Test
public void testRemoveFlatten() {
Node a = ParallelInstructionNode.create(new ParallelInstruction().setName("A"), Nodes.ExecutionLocation.UNKNOWN);
Node aPCollection = InstructionOutputNode.create(new InstructionOutput().setName("A.out"), PCOLLECTION_ID);
Edge aOutput = DefaultEdge.create();
Node b = ParallelInstructionNode.create(new ParallelInstruction().setName("B"), Nodes.ExecutionLocation.UNKNOWN);
Edge bOutput = DefaultEdge.create();
Node bPCollection = InstructionOutputNode.create(new InstructionOutput().setName("B.out"), PCOLLECTION_ID);
Node flatten = ParallelInstructionNode.create(new ParallelInstruction().setName("Flatten").setFlatten(new FlattenInstruction()), Nodes.ExecutionLocation.UNKNOWN);
Node flattenPCollection = InstructionOutputNode.create(new InstructionOutput().setName("Flatten.out"), PCOLLECTION_ID);
Node c = ParallelInstructionNode.create(new ParallelInstruction().setName("C"), Nodes.ExecutionLocation.UNKNOWN);
Edge cOutput = DefaultEdge.create();
Node cPCollection = InstructionOutputNode.create(new InstructionOutput().setName("C.out"), PCOLLECTION_ID);
// A --\
// Flatten --> C
// B --/
MutableNetwork<Node, Edge> network = createEmptyNetwork();
network.addNode(a);
network.addNode(aPCollection);
network.addNode(b);
network.addNode(bPCollection);
network.addNode(flatten);
network.addNode(flattenPCollection);
network.addNode(c);
network.addNode(cPCollection);
network.addEdge(a, aPCollection, aOutput);
network.addEdge(aPCollection, flatten, DefaultEdge.create());
network.addEdge(b, bPCollection, bOutput);
network.addEdge(bPCollection, flatten, DefaultEdge.create());
network.addEdge(flatten, flattenPCollection, DefaultEdge.create());
network.addEdge(flattenPCollection, c, DefaultEdge.create());
network.addEdge(c, cPCollection, cOutput);
// A --\
// C
// B --/
assertThatFlattenIsProperlyRemoved(network);
}
Aggregations