runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/graph/MapTaskToNetworkFunctionTest.java - beam - Git at Google

 /*
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
  * regarding copyright ownership.  The ASF licenses this file
  * to you under the Apache License, Version 2.0 (the
  * "License"); you may not use this file except in compliance
  * with the License.  You may obtain a copy of the License at
  *
  *     http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
 package org.apache.beam.runners.dataflow.worker.graph;

 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
 import static org.hamcrest.Matchers.instanceOf;
 import static org.hamcrest.collection.IsEmptyCollection.emptyCollectionOf;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThat;
 import static org.junit.Assert.assertTrue;

 import com.google.api.services.dataflow.model.FlattenInstruction;
 import com.google.api.services.dataflow.model.InstructionInput;
 import com.google.api.services.dataflow.model.InstructionOutput;
 import com.google.api.services.dataflow.model.MapTask;
 import com.google.api.services.dataflow.model.MultiOutputInfo;
 import com.google.api.services.dataflow.model.ParDoInstruction;
 import com.google.api.services.dataflow.model.ParallelInstruction;
 import com.google.api.services.dataflow.model.PartialGroupByKeyInstruction;
 import com.google.api.services.dataflow.model.ReadInstruction;
 import com.google.api.services.dataflow.model.WriteInstruction;
 import java.util.Arrays;
 import org.apache.beam.runners.dataflow.worker.graph.Edges.DefaultEdge;
 import org.apache.beam.runners.dataflow.worker.graph.Edges.Edge;
 import org.apache.beam.runners.dataflow.worker.graph.Edges.MultiOutputInfoEdge;
 import org.apache.beam.runners.dataflow.worker.graph.Nodes.InstructionOutputNode;
 import org.apache.beam.runners.dataflow.worker.graph.Nodes.Node;
 import org.apache.beam.runners.dataflow.worker.graph.Nodes.ParallelInstructionNode;
 import org.apache.beam.sdk.extensions.gcp.util.Transport;
 import org.apache.beam.sdk.fn.IdGenerators;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.graph.Network;
 import org.hamcrest.Matchers;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;

 /** Tests for {@link MapTaskToNetworkFunction}. */
 @RunWith(JUnit4.class)
 public class MapTaskToNetworkFunctionTest {
   @Test
   public void testEmptyMapTask() {
     Network<Node, Edge> network =
         new MapTaskToNetworkFunction(IdGenerators.decrementingLongs()).apply(new MapTask());
     assertTrue(network.isDirected());
     assertTrue(network.allowsParallelEdges());
     assertFalse(network.allowsSelfLoops());
     assertThat(network.nodes(), emptyCollectionOf(Node.class));
   }

   @Test
   public void testRead() {
     InstructionOutput readOutput = createInstructionOutput("Read.out");
     ParallelInstruction read = createParallelInstruction("Read", readOutput);
     read.setRead(new ReadInstruction());

     MapTask mapTask = new MapTask();
     mapTask.setInstructions(ImmutableList.of(read));
     mapTask.setFactory(Transport.getJsonFactory());

     Network<Node, Edge> network =
         new MapTaskToNetworkFunction(IdGenerators.decrementingLongs()).apply(mapTask);
     assertNetworkProperties(network);
     assertEquals(2, network.nodes().size());
     assertEquals(1, network.edges().size());

     ParallelInstructionNode readNode = get(network, read);
     InstructionOutputNode readOutputNode = getOnlySuccessor(network, readNode);
     assertEquals(readOutput, readOutputNode.getInstructionOutput());
   }

   @Test
   public void testParDo() {
     InstructionOutput readOutput = createInstructionOutput("Read.out");
     ParallelInstruction read = createParallelInstruction("Read", readOutput);
     read.setRead(new ReadInstruction());

     MultiOutputInfo parDoMultiOutput = createMultiOutputInfo("output");
     ParDoInstruction parDoInstruction = new ParDoInstruction();
     parDoInstruction.setInput(createInstructionInput(0, 0)); // Read.out
     parDoInstruction.setMultiOutputInfos(ImmutableList.of(parDoMultiOutput));
     InstructionOutput parDoOutput = createInstructionOutput("ParDo.out");
     ParallelInstruction parDo = createParallelInstruction("ParDo", parDoOutput);
     parDo.setParDo(parDoInstruction);

     MapTask mapTask = new MapTask();
     mapTask.setInstructions(ImmutableList.of(read, parDo));
     mapTask.setFactory(Transport.getJsonFactory());

     Network<Node, Edge> network =
         new MapTaskToNetworkFunction(IdGenerators.decrementingLongs()).apply(mapTask);
     assertNetworkProperties(network);
     assertEquals(4, network.nodes().size());
     assertEquals(3, network.edges().size());

     ParallelInstructionNode readNode = get(network, read);
     InstructionOutputNode readOutputNode = getOnlySuccessor(network, readNode);
     assertEquals(readOutput, readOutputNode.getInstructionOutput());

     ParallelInstructionNode parDoNode = getOnlySuccessor(network, readOutputNode);
     InstructionOutputNode parDoOutputNode = getOnlySuccessor(network, parDoNode);
     assertEquals(parDoOutput, parDoOutputNode.getInstructionOutput());

     assertEquals(
         parDoMultiOutput,
         ((MultiOutputInfoEdge)
                 Iterables.getOnlyElement(network.edgesConnecting(parDoNode, parDoOutputNode)))
             .getMultiOutputInfo());
   }

   @Test
   public void testFlatten() {
     // ReadA --\
     //          |--> Flatten
     // ReadB --/
     InstructionOutput readOutputA = createInstructionOutput("ReadA.out");
     ParallelInstruction readA = createParallelInstruction("ReadA", readOutputA);
     readA.setRead(new ReadInstruction());

     InstructionOutput readOutputB = createInstructionOutput("ReadB.out");
     ParallelInstruction readB = createParallelInstruction("ReadB", readOutputB);
     readB.setRead(new ReadInstruction());

     FlattenInstruction flattenInstruction = new FlattenInstruction();
     flattenInstruction.setInputs(
         ImmutableList.of(
             createInstructionInput(0, 0), // ReadA.out
             createInstructionInput(1, 0))); // ReadB.out
     InstructionOutput flattenOutput = createInstructionOutput("Flatten.out");
     ParallelInstruction flatten = createParallelInstruction("Flatten", flattenOutput);
     flatten.setFlatten(flattenInstruction);

     MapTask mapTask = new MapTask();
     mapTask.setInstructions(ImmutableList.of(readA, readB, flatten));
     mapTask.setFactory(Transport.getJsonFactory());

     Network<Node, Edge> network =
         new MapTaskToNetworkFunction(IdGenerators.decrementingLongs()).apply(mapTask);
     assertNetworkProperties(network);
     assertEquals(6, network.nodes().size());
     assertEquals(5, network.edges().size());

     ParallelInstructionNode readANode = get(network, readA);
     InstructionOutputNode readOutputANode = getOnlySuccessor(network, readANode);
     assertEquals(readOutputA, readOutputANode.getInstructionOutput());

     ParallelInstructionNode readBNode = get(network, readB);
     InstructionOutputNode readOutputBNode = getOnlySuccessor(network, readBNode);
     assertEquals(readOutputB, readOutputBNode.getInstructionOutput());

     // Make sure the successors for both ReadA and ReadB output PCollections are the same
     assertEquals(network.successors(readOutputANode), network.successors(readOutputBNode));

     ParallelInstructionNode flattenNode = getOnlySuccessor(network, readOutputANode);
     InstructionOutputNode flattenOutputNode = getOnlySuccessor(network, flattenNode);
     assertEquals(flattenOutput, flattenOutputNode.getInstructionOutput());
   }

   @Test
   public void testParallelEdgeFlatten() {
     //                  /---\
     // Read --> Read.out --> Flatten
     //                  \---/
     InstructionOutput readOutput = createInstructionOutput("Read.out");
     ParallelInstruction read = createParallelInstruction("Read", readOutput);
     read.setRead(new ReadInstruction());

     FlattenInstruction flattenInstruction = new FlattenInstruction();
     flattenInstruction.setInputs(
         ImmutableList.of(
             createInstructionInput(0, 0), // Read.out
             createInstructionInput(0, 0), // Read.out
             createInstructionInput(0, 0))); // Read.out
     InstructionOutput flattenOutput = createInstructionOutput("Flatten.out");
     ParallelInstruction flatten = createParallelInstruction("Flatten", flattenOutput);
     flatten.setFlatten(flattenInstruction);

     MapTask mapTask = new MapTask();
     mapTask.setInstructions(ImmutableList.of(read, flatten));
     mapTask.setFactory(Transport.getJsonFactory());

     Network<Node, Edge> network =
         new MapTaskToNetworkFunction(IdGenerators.decrementingLongs()).apply(mapTask);
     assertNetworkProperties(network);
     assertEquals(4, network.nodes().size());
     assertEquals(5, network.edges().size());

     ParallelInstructionNode readNode = get(network, read);
     InstructionOutputNode readOutputNode = getOnlySuccessor(network, readNode);
     assertEquals(readOutput, readOutputNode.getInstructionOutput());

     ParallelInstructionNode flattenNode = getOnlySuccessor(network, readOutputNode);
     // Assert that the three parallel edges are maintained
     assertEquals(3, network.edgesConnecting(readOutputNode, flattenNode).size());

     InstructionOutputNode flattenOutputNode = getOnlySuccessor(network, flattenNode);
     assertEquals(flattenOutput, flattenOutputNode.getInstructionOutput());
   }

   @Test
   public void testWrite() {
     InstructionOutput readOutput = createInstructionOutput("Read.out");
     ParallelInstruction read = createParallelInstruction("Read", readOutput);
     read.setRead(new ReadInstruction());

     WriteInstruction writeInstruction = new WriteInstruction();
     writeInstruction.setInput(createInstructionInput(0, 0)); // Read.out
     ParallelInstruction write = createParallelInstruction("Write");
     write.setWrite(writeInstruction);

     MapTask mapTask = new MapTask();
     mapTask.setInstructions(ImmutableList.of(read, write));
     mapTask.setFactory(Transport.getJsonFactory());

     Network<Node, Edge> network =
         new MapTaskToNetworkFunction(IdGenerators.decrementingLongs()).apply(mapTask);
     assertNetworkProperties(network);
     assertEquals(3, network.nodes().size());
     assertEquals(2, network.edges().size());

     ParallelInstructionNode readNode = get(network, read);
     InstructionOutputNode readOutputNode = getOnlySuccessor(network, readNode);
     assertEquals(readOutput, readOutputNode.getInstructionOutput());

     ParallelInstructionNode writeNode = getOnlySuccessor(network, readOutputNode);
     assertNotNull(writeNode);
   }

   @Test
   public void testPartialGroupByKey() {
     // Read --> PGBK --> Write
     InstructionOutput readOutput = createInstructionOutput("Read.out");
     ParallelInstruction read = createParallelInstruction("Read", readOutput);
     read.setRead(new ReadInstruction());

     PartialGroupByKeyInstruction pgbkInstruction = new PartialGroupByKeyInstruction();
     pgbkInstruction.setInput(createInstructionInput(0, 0)); // Read.out
     InstructionOutput pgbkOutput = createInstructionOutput("PGBK.out");
     ParallelInstruction pgbk = createParallelInstruction("PGBK", pgbkOutput);
     pgbk.setPartialGroupByKey(pgbkInstruction);

     WriteInstruction writeInstruction = new WriteInstruction();
     writeInstruction.setInput(createInstructionInput(1, 0)); // PGBK.out
     ParallelInstruction write = createParallelInstruction("Write");
     write.setWrite(writeInstruction);

     MapTask mapTask = new MapTask();
     mapTask.setInstructions(ImmutableList.of(read, pgbk, write));
     mapTask.setFactory(Transport.getJsonFactory());

     Network<Node, Edge> network =
         new MapTaskToNetworkFunction(IdGenerators.decrementingLongs()).apply(mapTask);
     assertNetworkProperties(network);
     assertEquals(5, network.nodes().size());
     assertEquals(4, network.edges().size());

     ParallelInstructionNode readNode = get(network, read);
     InstructionOutputNode readOutputNode = getOnlySuccessor(network, readNode);
     assertEquals(readOutput, readOutputNode.getInstructionOutput());

     ParallelInstructionNode pgbkNode = getOnlySuccessor(network, readOutputNode);
     InstructionOutputNode pgbkOutputNode = getOnlySuccessor(network, pgbkNode);
     assertEquals(pgbkOutput, pgbkOutputNode.getInstructionOutput());

     ParallelInstructionNode writeNode = getOnlySuccessor(network, pgbkOutputNode);
     assertNotNull(write);
   }

   @Test
   public void testMultipleOutput() {
     //                /---> WriteA
     // Read ---> ParDo
     //                \---> WriteB
     InstructionOutput readOutput = createInstructionOutput("Read.out");
     ParallelInstruction read = createParallelInstruction("Read", readOutput);
     read.setRead(new ReadInstruction());

     MultiOutputInfo parDoMultiOutput1 = createMultiOutputInfo("output1");
     MultiOutputInfo parDoMultiOutput2 = createMultiOutputInfo("output2");
     ParDoInstruction parDoInstruction = new ParDoInstruction();
     parDoInstruction.setInput(createInstructionInput(0, 0)); // Read.out
     parDoInstruction.setMultiOutputInfos(ImmutableList.of(parDoMultiOutput1, parDoMultiOutput2));
     InstructionOutput parDoOutput1 = createInstructionOutput("ParDo.out1");
     InstructionOutput parDoOutput2 = createInstructionOutput("ParDo.out2");
     ParallelInstruction parDo = createParallelInstruction("ParDo", parDoOutput1, parDoOutput2);
     parDo.setParDo(parDoInstruction);

     WriteInstruction writeAInstruction = new WriteInstruction();
     writeAInstruction.setInput(createInstructionInput(1, 0)); // ParDo.out1
     ParallelInstruction writeA = createParallelInstruction("WriteA");
     writeA.setWrite(writeAInstruction);

     WriteInstruction writeBInstruction = new WriteInstruction();
     writeBInstruction.setInput(createInstructionInput(1, 1)); // ParDo.out2
     ParallelInstruction writeB = createParallelInstruction("WriteB");
     writeB.setWrite(writeBInstruction);

     MapTask mapTask = new MapTask();
     mapTask.setInstructions(ImmutableList.of(read, parDo, writeA, writeB));
     mapTask.setFactory(Transport.getJsonFactory());

     Network<Node, Edge> network =
         new MapTaskToNetworkFunction(IdGenerators.decrementingLongs()).apply(mapTask);
     assertNetworkProperties(network);
     assertEquals(7, network.nodes().size());
     assertEquals(6, network.edges().size());

     ParallelInstructionNode parDoNode = get(network, parDo);
     ParallelInstructionNode writeANode = get(network, writeA);
     ParallelInstructionNode writeBNode = get(network, writeB);
     InstructionOutputNode parDoOutput1Node = getOnlyPredecessor(network, writeANode);
     assertEquals(parDoOutput1, parDoOutput1Node.getInstructionOutput());

     InstructionOutputNode parDoOutput2Node = getOnlyPredecessor(network, writeBNode);
     assertEquals(parDoOutput2, parDoOutput2Node.getInstructionOutput());

     assertThat(
         network.successors(parDoNode),
         Matchers.<Node>containsInAnyOrder(parDoOutput1Node, parDoOutput2Node));

     assertEquals(
         parDoMultiOutput1,
         ((MultiOutputInfoEdge)
                 Iterables.getOnlyElement(network.edgesConnecting(parDoNode, parDoOutput1Node)))
             .getMultiOutputInfo());
     assertEquals(
         parDoMultiOutput2,
         ((MultiOutputInfoEdge)
                 Iterables.getOnlyElement(network.edgesConnecting(parDoNode, parDoOutput2Node)))
             .getMultiOutputInfo());
   }

   private static void assertNetworkProperties(Network<Node, Edge> network) {
     assertTrue(network.isDirected());
     assertFalse(network.allowsSelfLoops());
     for (Node node : network.nodes()) {
       if (node instanceof ParallelInstructionNode) {
         assertNull(((ParallelInstructionNode) node).getParallelInstruction().getOutputs());
         // Validate that all successors are PCollections with DefaultEdge outgoing edges
         // except for ParDoInstructions
         for (Node successor : network.successors(node)) {
           assertThat(successor, instanceOf(InstructionOutputNode.class));
           // Assert that all outgoing edges for a ParDo are MultiOutputInfoEdges
           if (((ParallelInstructionNode) node).getParallelInstruction().getParDo() != null) {
             for (Edge edge : network.edgesConnecting(node, successor)) {
               assertThat(edge, instanceOf(MultiOutputInfoEdge.class));
             }
           } else {
             for (Edge edge : network.edgesConnecting(node, successor)) {
               assertThat(edge, instanceOf(DefaultEdge.class));
             }
           }
         }

       } else if (node instanceof InstructionOutputNode) {
         assertThat(network.inDegree(node), greaterThanOrEqualTo(1));
         // Validate that all successors are instructions with DefaultEdge outgoing edges
         for (Node successor : network.successors(node)) {
           assertThat(successor, instanceOf(ParallelInstructionNode.class));
           for (Edge edge : network.edgesConnecting(node, successor)) {
             assertThat(edge, instanceOf(DefaultEdge.class));
           }
         }
       }
     }
   }

   private static ParallelInstructionNode get(
       Network<Node, Edge> network, ParallelInstruction instruction) {
     for (Node node : network.nodes()) {
       if (node instanceof ParallelInstructionNode
           && ((ParallelInstructionNode) node)
               .getParallelInstruction()
               .getName()
               .equals(instruction.getName())) {
         return (ParallelInstructionNode) node;
       }
     }
     throw new AssertionError(
         String.format("Unable to find instruction %s in network %s", instruction, network));
   }

   private static ParallelInstructionNode getOnlySuccessor(
       Network<Node, Edge> network, InstructionOutputNode node) {
     return (ParallelInstructionNode) Iterables.getOnlyElement(network.successors(node));
   }

   private static InstructionOutputNode getOnlySuccessor(
       Network<Node, Edge> network, ParallelInstructionNode node) {
     return (InstructionOutputNode) Iterables.getOnlyElement(network.successors(node));
   }

   private static InstructionOutputNode getOnlyPredecessor(
       Network<Node, Edge> network, ParallelInstructionNode node) {
     return (InstructionOutputNode) Iterables.getOnlyElement(network.predecessors(node));
   }

   private static InstructionOutput createInstructionOutput(String name) {
     InstructionOutput rval = new InstructionOutput();
     rval.setName(name);
     return rval;
   }

   private static ParallelInstruction createParallelInstruction(
       String name, InstructionOutput... outputs) {
     ParallelInstruction rval = new ParallelInstruction();
     rval.setName(name);
     rval.setOutputs(Arrays.asList(outputs));
     return rval;
   }

   private static MultiOutputInfo createMultiOutputInfo(String tag) {
     MultiOutputInfo rval = new MultiOutputInfo();
     rval.setTag(tag);
     return rval;
   }

   private static InstructionInput createInstructionInput(int instructionIndex, int outputNum) {
     InstructionInput rval = new InstructionInput();
     rval.setProducerInstructionIndex(instructionIndex);
     rval.setOutputNum(outputNum);
     return rval;
   }
 }
	/*
	* Licensed to the Apache Software Foundation (ASF) under one
	* or more contributor license agreements. See the NOTICE file
	* distributed with this work for additional information
	* regarding copyright ownership. The ASF licenses this file
	* to you under the Apache License, Version 2.0 (the
	* "License"); you may not use this file except in compliance
	* with the License. You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing, software
	* distributed under the License is distributed on an "AS IS" BASIS,
	* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	* See the License for the specific language governing permissions and
	* limitations under the License.
	*/
	package org.apache.beam.runners.dataflow.worker.graph;

	import static org.hamcrest.Matchers.greaterThanOrEqualTo;
	import static org.hamcrest.Matchers.instanceOf;
	import static org.hamcrest.collection.IsEmptyCollection.emptyCollectionOf;
	import static org.junit.Assert.assertEquals;
	import static org.junit.Assert.assertFalse;
	import static org.junit.Assert.assertNotNull;
	import static org.junit.Assert.assertNull;
	import static org.junit.Assert.assertThat;
	import static org.junit.Assert.assertTrue;

	import com.google.api.services.dataflow.model.FlattenInstruction;
	import com.google.api.services.dataflow.model.InstructionInput;
	import com.google.api.services.dataflow.model.InstructionOutput;
	import com.google.api.services.dataflow.model.MapTask;
	import com.google.api.services.dataflow.model.MultiOutputInfo;
	import com.google.api.services.dataflow.model.ParDoInstruction;
	import com.google.api.services.dataflow.model.ParallelInstruction;
	import com.google.api.services.dataflow.model.PartialGroupByKeyInstruction;
	import com.google.api.services.dataflow.model.ReadInstruction;
	import com.google.api.services.dataflow.model.WriteInstruction;
	import java.util.Arrays;
	import org.apache.beam.runners.dataflow.worker.graph.Edges.DefaultEdge;
	import org.apache.beam.runners.dataflow.worker.graph.Edges.Edge;
	import org.apache.beam.runners.dataflow.worker.graph.Edges.MultiOutputInfoEdge;
	import org.apache.beam.runners.dataflow.worker.graph.Nodes.InstructionOutputNode;
	import org.apache.beam.runners.dataflow.worker.graph.Nodes.Node;
	import org.apache.beam.runners.dataflow.worker.graph.Nodes.ParallelInstructionNode;
	import org.apache.beam.sdk.extensions.gcp.util.Transport;
	import org.apache.beam.sdk.fn.IdGenerators;
	import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
	import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
	import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.graph.Network;
	import org.hamcrest.Matchers;
	import org.junit.Test;
	import org.junit.runner.RunWith;
	import org.junit.runners.JUnit4;

	/** Tests for {@link MapTaskToNetworkFunction}. */
	@RunWith(JUnit4.class)
	public class MapTaskToNetworkFunctionTest {
	@Test
	public void testEmptyMapTask() {
	Network<Node, Edge> network =
	new MapTaskToNetworkFunction(IdGenerators.decrementingLongs()).apply(new MapTask());
	assertTrue(network.isDirected());
	assertTrue(network.allowsParallelEdges());
	assertFalse(network.allowsSelfLoops());
	assertThat(network.nodes(), emptyCollectionOf(Node.class));
	}

	@Test
	public void testRead() {
	InstructionOutput readOutput = createInstructionOutput("Read.out");
	ParallelInstruction read = createParallelInstruction("Read", readOutput);
	read.setRead(new ReadInstruction());

	MapTask mapTask = new MapTask();
	mapTask.setInstructions(ImmutableList.of(read));
	mapTask.setFactory(Transport.getJsonFactory());

	Network<Node, Edge> network =
	new MapTaskToNetworkFunction(IdGenerators.decrementingLongs()).apply(mapTask);
	assertNetworkProperties(network);
	assertEquals(2, network.nodes().size());
	assertEquals(1, network.edges().size());

	ParallelInstructionNode readNode = get(network, read);
	InstructionOutputNode readOutputNode = getOnlySuccessor(network, readNode);
	assertEquals(readOutput, readOutputNode.getInstructionOutput());
	}

	@Test
	public void testParDo() {
	InstructionOutput readOutput = createInstructionOutput("Read.out");
	ParallelInstruction read = createParallelInstruction("Read", readOutput);
	read.setRead(new ReadInstruction());

	MultiOutputInfo parDoMultiOutput = createMultiOutputInfo("output");
	ParDoInstruction parDoInstruction = new ParDoInstruction();
	parDoInstruction.setInput(createInstructionInput(0, 0)); // Read.out
	parDoInstruction.setMultiOutputInfos(ImmutableList.of(parDoMultiOutput));
	InstructionOutput parDoOutput = createInstructionOutput("ParDo.out");
	ParallelInstruction parDo = createParallelInstruction("ParDo", parDoOutput);
	parDo.setParDo(parDoInstruction);

	MapTask mapTask = new MapTask();
	mapTask.setInstructions(ImmutableList.of(read, parDo));
	mapTask.setFactory(Transport.getJsonFactory());

	Network<Node, Edge> network =
	new MapTaskToNetworkFunction(IdGenerators.decrementingLongs()).apply(mapTask);
	assertNetworkProperties(network);
	assertEquals(4, network.nodes().size());
	assertEquals(3, network.edges().size());

	ParallelInstructionNode readNode = get(network, read);
	InstructionOutputNode readOutputNode = getOnlySuccessor(network, readNode);
	assertEquals(readOutput, readOutputNode.getInstructionOutput());

	ParallelInstructionNode parDoNode = getOnlySuccessor(network, readOutputNode);
	InstructionOutputNode parDoOutputNode = getOnlySuccessor(network, parDoNode);
	assertEquals(parDoOutput, parDoOutputNode.getInstructionOutput());

	assertEquals(
	parDoMultiOutput,
	((MultiOutputInfoEdge)
	Iterables.getOnlyElement(network.edgesConnecting(parDoNode, parDoOutputNode)))
	.getMultiOutputInfo());
	}

	@Test
	public void testFlatten() {
	// ReadA --\
	// \|--> Flatten
	// ReadB --/
	InstructionOutput readOutputA = createInstructionOutput("ReadA.out");
	ParallelInstruction readA = createParallelInstruction("ReadA", readOutputA);
	readA.setRead(new ReadInstruction());

	InstructionOutput readOutputB = createInstructionOutput("ReadB.out");
	ParallelInstruction readB = createParallelInstruction("ReadB", readOutputB);
	readB.setRead(new ReadInstruction());

	FlattenInstruction flattenInstruction = new FlattenInstruction();
	flattenInstruction.setInputs(
	ImmutableList.of(
	createInstructionInput(0, 0), // ReadA.out
	createInstructionInput(1, 0))); // ReadB.out
	InstructionOutput flattenOutput = createInstructionOutput("Flatten.out");
	ParallelInstruction flatten = createParallelInstruction("Flatten", flattenOutput);
	flatten.setFlatten(flattenInstruction);

	MapTask mapTask = new MapTask();
	mapTask.setInstructions(ImmutableList.of(readA, readB, flatten));
	mapTask.setFactory(Transport.getJsonFactory());

	Network<Node, Edge> network =
	new MapTaskToNetworkFunction(IdGenerators.decrementingLongs()).apply(mapTask);
	assertNetworkProperties(network);
	assertEquals(6, network.nodes().size());
	assertEquals(5, network.edges().size());

	ParallelInstructionNode readANode = get(network, readA);
	InstructionOutputNode readOutputANode = getOnlySuccessor(network, readANode);
	assertEquals(readOutputA, readOutputANode.getInstructionOutput());

	ParallelInstructionNode readBNode = get(network, readB);
	InstructionOutputNode readOutputBNode = getOnlySuccessor(network, readBNode);
	assertEquals(readOutputB, readOutputBNode.getInstructionOutput());

	// Make sure the successors for both ReadA and ReadB output PCollections are the same
	assertEquals(network.successors(readOutputANode), network.successors(readOutputBNode));

	ParallelInstructionNode flattenNode = getOnlySuccessor(network, readOutputANode);
	InstructionOutputNode flattenOutputNode = getOnlySuccessor(network, flattenNode);
	assertEquals(flattenOutput, flattenOutputNode.getInstructionOutput());
	}

	@Test
	public void testParallelEdgeFlatten() {
	// /---\
	// Read --> Read.out --> Flatten
	// \---/
	InstructionOutput readOutput = createInstructionOutput("Read.out");
	ParallelInstruction read = createParallelInstruction("Read", readOutput);
	read.setRead(new ReadInstruction());

	FlattenInstruction flattenInstruction = new FlattenInstruction();
	flattenInstruction.setInputs(
	ImmutableList.of(
	createInstructionInput(0, 0), // Read.out
	createInstructionInput(0, 0), // Read.out
	createInstructionInput(0, 0))); // Read.out
	InstructionOutput flattenOutput = createInstructionOutput("Flatten.out");
	ParallelInstruction flatten = createParallelInstruction("Flatten", flattenOutput);
	flatten.setFlatten(flattenInstruction);

	MapTask mapTask = new MapTask();
	mapTask.setInstructions(ImmutableList.of(read, flatten));
	mapTask.setFactory(Transport.getJsonFactory());

	Network<Node, Edge> network =
	new MapTaskToNetworkFunction(IdGenerators.decrementingLongs()).apply(mapTask);
	assertNetworkProperties(network);
	assertEquals(4, network.nodes().size());
	assertEquals(5, network.edges().size());

	ParallelInstructionNode readNode = get(network, read);
	InstructionOutputNode readOutputNode = getOnlySuccessor(network, readNode);
	assertEquals(readOutput, readOutputNode.getInstructionOutput());

	ParallelInstructionNode flattenNode = getOnlySuccessor(network, readOutputNode);
	// Assert that the three parallel edges are maintained
	assertEquals(3, network.edgesConnecting(readOutputNode, flattenNode).size());

	InstructionOutputNode flattenOutputNode = getOnlySuccessor(network, flattenNode);
	assertEquals(flattenOutput, flattenOutputNode.getInstructionOutput());
	}

	@Test
	public void testWrite() {
	InstructionOutput readOutput = createInstructionOutput("Read.out");
	ParallelInstruction read = createParallelInstruction("Read", readOutput);
	read.setRead(new ReadInstruction());

	WriteInstruction writeInstruction = new WriteInstruction();
	writeInstruction.setInput(createInstructionInput(0, 0)); // Read.out
	ParallelInstruction write = createParallelInstruction("Write");
	write.setWrite(writeInstruction);

	MapTask mapTask = new MapTask();
	mapTask.setInstructions(ImmutableList.of(read, write));
	mapTask.setFactory(Transport.getJsonFactory());

	Network<Node, Edge> network =
	new MapTaskToNetworkFunction(IdGenerators.decrementingLongs()).apply(mapTask);
	assertNetworkProperties(network);
	assertEquals(3, network.nodes().size());
	assertEquals(2, network.edges().size());

	ParallelInstructionNode readNode = get(network, read);
	InstructionOutputNode readOutputNode = getOnlySuccessor(network, readNode);
	assertEquals(readOutput, readOutputNode.getInstructionOutput());

	ParallelInstructionNode writeNode = getOnlySuccessor(network, readOutputNode);
	assertNotNull(writeNode);
	}

	@Test
	public void testPartialGroupByKey() {
	// Read --> PGBK --> Write
	InstructionOutput readOutput = createInstructionOutput("Read.out");
	ParallelInstruction read = createParallelInstruction("Read", readOutput);
	read.setRead(new ReadInstruction());

	PartialGroupByKeyInstruction pgbkInstruction = new PartialGroupByKeyInstruction();
	pgbkInstruction.setInput(createInstructionInput(0, 0)); // Read.out
	InstructionOutput pgbkOutput = createInstructionOutput("PGBK.out");
	ParallelInstruction pgbk = createParallelInstruction("PGBK", pgbkOutput);
	pgbk.setPartialGroupByKey(pgbkInstruction);

	WriteInstruction writeInstruction = new WriteInstruction();
	writeInstruction.setInput(createInstructionInput(1, 0)); // PGBK.out
	ParallelInstruction write = createParallelInstruction("Write");
	write.setWrite(writeInstruction);

	MapTask mapTask = new MapTask();
	mapTask.setInstructions(ImmutableList.of(read, pgbk, write));
	mapTask.setFactory(Transport.getJsonFactory());

	Network<Node, Edge> network =
	new MapTaskToNetworkFunction(IdGenerators.decrementingLongs()).apply(mapTask);
	assertNetworkProperties(network);
	assertEquals(5, network.nodes().size());
	assertEquals(4, network.edges().size());

	ParallelInstructionNode readNode = get(network, read);
	InstructionOutputNode readOutputNode = getOnlySuccessor(network, readNode);
	assertEquals(readOutput, readOutputNode.getInstructionOutput());

	ParallelInstructionNode pgbkNode = getOnlySuccessor(network, readOutputNode);
	InstructionOutputNode pgbkOutputNode = getOnlySuccessor(network, pgbkNode);
	assertEquals(pgbkOutput, pgbkOutputNode.getInstructionOutput());

	ParallelInstructionNode writeNode = getOnlySuccessor(network, pgbkOutputNode);
	assertNotNull(write);
	}

	@Test
	public void testMultipleOutput() {
	// /---> WriteA
	// Read ---> ParDo
	// \---> WriteB
	InstructionOutput readOutput = createInstructionOutput("Read.out");
	ParallelInstruction read = createParallelInstruction("Read", readOutput);
	read.setRead(new ReadInstruction());

	MultiOutputInfo parDoMultiOutput1 = createMultiOutputInfo("output1");
	MultiOutputInfo parDoMultiOutput2 = createMultiOutputInfo("output2");
	ParDoInstruction parDoInstruction = new ParDoInstruction();
	parDoInstruction.setInput(createInstructionInput(0, 0)); // Read.out
	parDoInstruction.setMultiOutputInfos(ImmutableList.of(parDoMultiOutput1, parDoMultiOutput2));
	InstructionOutput parDoOutput1 = createInstructionOutput("ParDo.out1");
	InstructionOutput parDoOutput2 = createInstructionOutput("ParDo.out2");
	ParallelInstruction parDo = createParallelInstruction("ParDo", parDoOutput1, parDoOutput2);
	parDo.setParDo(parDoInstruction);

	WriteInstruction writeAInstruction = new WriteInstruction();
	writeAInstruction.setInput(createInstructionInput(1, 0)); // ParDo.out1
	ParallelInstruction writeA = createParallelInstruction("WriteA");
	writeA.setWrite(writeAInstruction);

	WriteInstruction writeBInstruction = new WriteInstruction();
	writeBInstruction.setInput(createInstructionInput(1, 1)); // ParDo.out2
	ParallelInstruction writeB = createParallelInstruction("WriteB");
	writeB.setWrite(writeBInstruction);

	MapTask mapTask = new MapTask();
	mapTask.setInstructions(ImmutableList.of(read, parDo, writeA, writeB));
	mapTask.setFactory(Transport.getJsonFactory());

	Network<Node, Edge> network =
	new MapTaskToNetworkFunction(IdGenerators.decrementingLongs()).apply(mapTask);
	assertNetworkProperties(network);
	assertEquals(7, network.nodes().size());
	assertEquals(6, network.edges().size());

	ParallelInstructionNode parDoNode = get(network, parDo);
	ParallelInstructionNode writeANode = get(network, writeA);
	ParallelInstructionNode writeBNode = get(network, writeB);
	InstructionOutputNode parDoOutput1Node = getOnlyPredecessor(network, writeANode);
	assertEquals(parDoOutput1, parDoOutput1Node.getInstructionOutput());

	InstructionOutputNode parDoOutput2Node = getOnlyPredecessor(network, writeBNode);
	assertEquals(parDoOutput2, parDoOutput2Node.getInstructionOutput());

	assertThat(
	network.successors(parDoNode),
	Matchers.<Node>containsInAnyOrder(parDoOutput1Node, parDoOutput2Node));

	assertEquals(
	parDoMultiOutput1,
	((MultiOutputInfoEdge)
	Iterables.getOnlyElement(network.edgesConnecting(parDoNode, parDoOutput1Node)))
	.getMultiOutputInfo());
	assertEquals(
	parDoMultiOutput2,
	((MultiOutputInfoEdge)
	Iterables.getOnlyElement(network.edgesConnecting(parDoNode, parDoOutput2Node)))
	.getMultiOutputInfo());
	}

	private static void assertNetworkProperties(Network<Node, Edge> network) {
	assertTrue(network.isDirected());
	assertFalse(network.allowsSelfLoops());
	for (Node node : network.nodes()) {
	if (node instanceof ParallelInstructionNode) {
	assertNull(((ParallelInstructionNode) node).getParallelInstruction().getOutputs());
	// Validate that all successors are PCollections with DefaultEdge outgoing edges
	// except for ParDoInstructions
	for (Node successor : network.successors(node)) {
	assertThat(successor, instanceOf(InstructionOutputNode.class));
	// Assert that all outgoing edges for a ParDo are MultiOutputInfoEdges
	if (((ParallelInstructionNode) node).getParallelInstruction().getParDo() != null) {
	for (Edge edge : network.edgesConnecting(node, successor)) {
	assertThat(edge, instanceOf(MultiOutputInfoEdge.class));
	}
	} else {
	for (Edge edge : network.edgesConnecting(node, successor)) {
	assertThat(edge, instanceOf(DefaultEdge.class));
	}
	}
	}

	} else if (node instanceof InstructionOutputNode) {
	assertThat(network.inDegree(node), greaterThanOrEqualTo(1));
	// Validate that all successors are instructions with DefaultEdge outgoing edges
	for (Node successor : network.successors(node)) {
	assertThat(successor, instanceOf(ParallelInstructionNode.class));
	for (Edge edge : network.edgesConnecting(node, successor)) {
	assertThat(edge, instanceOf(DefaultEdge.class));
	}
	}
	}
	}
	}

	private static ParallelInstructionNode get(
	Network<Node, Edge> network, ParallelInstruction instruction) {
	for (Node node : network.nodes()) {
	if (node instanceof ParallelInstructionNode
	&& ((ParallelInstructionNode) node)
	.getParallelInstruction()
	.getName()
	.equals(instruction.getName())) {
	return (ParallelInstructionNode) node;
	}
	}
	throw new AssertionError(
	String.format("Unable to find instruction %s in network %s", instruction, network));
	}

	private static ParallelInstructionNode getOnlySuccessor(
	Network<Node, Edge> network, InstructionOutputNode node) {
	return (ParallelInstructionNode) Iterables.getOnlyElement(network.successors(node));
	}

	private static InstructionOutputNode getOnlySuccessor(
	Network<Node, Edge> network, ParallelInstructionNode node) {
	return (InstructionOutputNode) Iterables.getOnlyElement(network.successors(node));
	}

	private static InstructionOutputNode getOnlyPredecessor(
	Network<Node, Edge> network, ParallelInstructionNode node) {
	return (InstructionOutputNode) Iterables.getOnlyElement(network.predecessors(node));
	}

	private static InstructionOutput createInstructionOutput(String name) {
	InstructionOutput rval = new InstructionOutput();
	rval.setName(name);
	return rval;
	}

	private static ParallelInstruction createParallelInstruction(
	String name, InstructionOutput... outputs) {
	ParallelInstruction rval = new ParallelInstruction();
	rval.setName(name);
	rval.setOutputs(Arrays.asList(outputs));
	return rval;
	}

	private static MultiOutputInfo createMultiOutputInfo(String tag) {
	MultiOutputInfo rval = new MultiOutputInfo();
	rval.setTag(tag);
	return rval;
	}

	private static InstructionInput createInstructionInput(int instructionIndex, int outputNum) {
	InstructionInput rval = new InstructionInput();
	rval.setProducerInstructionIndex(instructionIndex);
	rval.setOutputNum(outputNum);
	return rval;
	}
	}