/*
 * 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.containsInAnyOrder;
import static org.hamcrest.Matchers.everyItem;
import static org.hamcrest.Matchers.instanceOf;
import static org.hamcrest.Matchers.not;
import static org.junit.Assert.assertThat;
import static org.junit.Assert.assertTrue;
import static org.mockito.Matchers.anyString;
import static org.mockito.Mockito.when;

import com.google.api.services.dataflow.model.InstructionOutput;
import com.google.api.services.dataflow.model.ParDoInstruction;
import com.google.api.services.dataflow.model.ParallelInstruction;
import com.google.api.services.dataflow.model.ReadInstruction;
import com.google.auto.value.AutoValue;
import java.util.List;
import java.util.function.BiFunction;
import java.util.function.Function;
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.HappensBeforeEdge;
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.fn.IdGenerators;
import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Iterables;
import org.apache.beam.vendor.guava.v20_0.com.google.common.graph.Graphs;
import org.apache.beam.vendor.guava.v20_0.com.google.common.graph.MutableNetwork;
import org.apache.beam.vendor.guava.v20_0.com.google.common.graph.Network;
import org.apache.beam.vendor.guava.v20_0.com.google.common.graph.NetworkBuilder;
import org.hamcrest.Matchers;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
import org.mockito.ArgumentCaptor;
import org.mockito.Mock;
import org.mockito.MockitoAnnotations;

/** Tests for {@link CreateRegisterFnOperationFunction}. */
@RunWith(JUnit4.class)
public class CreateRegisterFnOperationFunctionTest {

  @Mock private BiFunction<String, String, Node> portSupplier;
  @Mock private Function<MutableNetwork<Node, Edge>, Node> registerFnOperationFunction;
  private Function<MutableNetwork<Node, Edge>, MutableNetwork<Node, Edge>>
      createRegisterFnOperation;

  @Before
  public void setUp() {
    MockitoAnnotations.initMocks(this);
    createRegisterFnOperation =
        new CreateRegisterFnOperationFunction(
            IdGenerators.decrementingLongs(), portSupplier, registerFnOperationFunction, false);
  }

  @Test
  public void testEmptyGraph() {
    MutableNetwork<Node, Edge> appliedNetwork =
        createRegisterFnOperation.apply(createEmptyNetwork());
    MutableNetwork<Node, Edge> expectedNetwork = createEmptyNetwork();

    assertNetworkMaintainsBipartiteStructure(appliedNetwork);
    assertTrue(
        String.format("Expected network %s but got network %s", expectedNetwork, appliedNetwork),
        Graphs.equivalent(expectedNetwork, appliedNetwork));
  }

  @Test
  public void testAllRunnerGraph() {
    Node readNode = createReadNode("Read", Nodes.ExecutionLocation.RUNNER_HARNESS);
    Edge readNodeEdge = DefaultEdge.create();
    Node readNodeOut = createInstructionOutputNode("Read.out");
    Edge readNodeOutEdge = DefaultEdge.create();
    Node parDoNode = createParDoNode("ParDo", Nodes.ExecutionLocation.RUNNER_HARNESS);
    Edge parDoNodeEdge = DefaultEdge.create();
    Node parDoNodeOut = createInstructionOutputNode("ParDo.out");

    // Read -out-> ParDo
    MutableNetwork<Node, Edge> expectedNetwork = createEmptyNetwork();
    expectedNetwork.addNode(readNode);
    expectedNetwork.addNode(readNodeOut);
    expectedNetwork.addNode(parDoNode);
    expectedNetwork.addNode(parDoNodeOut);
    expectedNetwork.addEdge(readNode, readNodeOut, readNodeEdge);
    expectedNetwork.addEdge(readNodeOut, parDoNode, readNodeOutEdge);
    expectedNetwork.addEdge(parDoNode, parDoNodeOut, parDoNodeEdge);

    MutableNetwork<Node, Edge> appliedNetwork =
        createRegisterFnOperation.apply(Graphs.copyOf(expectedNetwork));

    assertNetworkMaintainsBipartiteStructure(appliedNetwork);
    assertTrue(
        String.format("Expected network %s but got network %s", expectedNetwork, appliedNetwork),
        Graphs.equivalent(expectedNetwork, appliedNetwork));
  }

  @Test
  public void testAllSdkGraph() {
    Node sdkPortionNode = TestNode.create("SdkPortion");
    @SuppressWarnings({"unchecked", "rawtypes"})
    ArgumentCaptor<MutableNetwork<Node, Edge>> networkCapture =
        ArgumentCaptor.forClass((Class) MutableNetwork.class);

    when(registerFnOperationFunction.apply(networkCapture.capture())).thenReturn(sdkPortionNode);

    // Read -out-> ParDo
    Node readNode = createReadNode("Read", Nodes.ExecutionLocation.SDK_HARNESS);
    Edge readNodeEdge = DefaultEdge.create();
    Node readNodeOut = createInstructionOutputNode("Read.out");
    Edge readNodeOutEdge = DefaultEdge.create();
    Node parDoNode = createParDoNode("ParDo", Nodes.ExecutionLocation.SDK_HARNESS);
    Edge parDoNodeEdge = DefaultEdge.create();
    Node parDoNodeOut = createInstructionOutputNode("ParDo.out");

    MutableNetwork<Node, Edge> network = createEmptyNetwork();
    network.addNode(readNode);
    network.addNode(readNodeOut);
    network.addNode(parDoNode);
    network.addNode(parDoNodeOut);
    network.addEdge(readNode, readNodeOut, readNodeEdge);
    network.addEdge(readNodeOut, parDoNode, readNodeOutEdge);
    network.addEdge(parDoNode, parDoNodeOut, parDoNodeEdge);

    MutableNetwork<Node, Edge> expectedNetwork = createEmptyNetwork();
    expectedNetwork.addNode(sdkPortionNode);

    MutableNetwork<Node, Edge> appliedNetwork =
        createRegisterFnOperation.apply(Graphs.copyOf(network));

    assertNetworkMaintainsBipartiteStructure(appliedNetwork);
    assertNetworkMaintainsBipartiteStructure(networkCapture.getValue());
    assertTrue(
        String.format("Expected network %s but got network %s", expectedNetwork, appliedNetwork),
        Graphs.equivalent(expectedNetwork, appliedNetwork));
    assertTrue(
        String.format("Expected network %s but got network %s", network, networkCapture.getValue()),
        Graphs.equivalent(network, networkCapture.getValue()));
  }

  @Test
  public void testRunnerToSdkToRunnerGraph() {
    Node sdkPortion = TestNode.create("SdkPortion");
    @SuppressWarnings({"unchecked", "rawtypes"})
    ArgumentCaptor<MutableNetwork<Node, Edge>> networkCapture =
        ArgumentCaptor.forClass((Class) MutableNetwork.class);
    when(registerFnOperationFunction.apply(networkCapture.capture())).thenReturn(sdkPortion);

    Node firstPort = TestNode.create("FirstPort");
    Node secondPort = TestNode.create("SecondPort");
    when(portSupplier.apply(anyString(), anyString())).thenReturn(firstPort, secondPort);

    Node readNode = createReadNode("Read", Nodes.ExecutionLocation.RUNNER_HARNESS);
    Edge readNodeEdge = DefaultEdge.create();
    Node readNodeOut = createInstructionOutputNode("Read.out");
    Edge readNodeOutEdge = DefaultEdge.create();
    Node sdkParDoNode = createParDoNode("SdkParDo", Nodes.ExecutionLocation.SDK_HARNESS);
    Edge sdkParDoNodeEdge = DefaultEdge.create();
    Node sdkParDoNodeOut = createInstructionOutputNode("SdkParDo.out");
    Edge sdkParDoNodeOutEdge = DefaultEdge.create();
    Node runnerParDoNode = createParDoNode("RunnerParDo", Nodes.ExecutionLocation.RUNNER_HARNESS);
    Edge runnerParDoNodeEdge = DefaultEdge.create();
    Node runnerParDoNodeOut = createInstructionOutputNode("RunnerParDo.out");

    // Read -out-> SdkParDo -out-> RunnerParDo
    MutableNetwork<Node, Edge> network = createEmptyNetwork();
    network.addNode(readNode);
    network.addNode(readNodeOut);
    network.addNode(sdkParDoNodeOut);
    network.addNode(sdkParDoNodeOut);
    network.addNode(runnerParDoNode);
    network.addNode(runnerParDoNodeOut);
    network.addEdge(readNode, readNodeOut, readNodeEdge);
    network.addEdge(readNodeOut, sdkParDoNode, readNodeOutEdge);
    network.addEdge(sdkParDoNode, sdkParDoNodeOut, sdkParDoNodeEdge);
    network.addEdge(sdkParDoNodeOut, runnerParDoNode, sdkParDoNodeOutEdge);
    network.addEdge(runnerParDoNode, runnerParDoNodeOut, runnerParDoNodeEdge);

    MutableNetwork<Node, Edge> appliedNetwork =
        createRegisterFnOperation.apply(Graphs.copyOf(network));
    assertNetworkMaintainsBipartiteStructure(appliedNetwork);

    // On each rewire between runner and SDK and vice versa, we use a new output node
    Node newOutA = Iterables.getOnlyElement(appliedNetwork.predecessors(firstPort));
    Node newOutB = Iterables.getOnlyElement(appliedNetwork.successors(secondPort));

    // readNode -newOutA-> firstPort --> sdkPortion --> secondPort -newOutB-> runnerParDoNode
    assertThat(
        appliedNetwork.nodes(),
        containsInAnyOrder(
            readNode,
            newOutA,
            firstPort,
            sdkPortion,
            secondPort,
            newOutB,
            runnerParDoNode,
            runnerParDoNodeOut));
    assertThat(appliedNetwork.successors(readNode), containsInAnyOrder(newOutA));
    assertThat(appliedNetwork.successors(newOutA), containsInAnyOrder(firstPort));
    assertThat(appliedNetwork.successors(firstPort), containsInAnyOrder(sdkPortion));
    assertThat(appliedNetwork.successors(sdkPortion), containsInAnyOrder(secondPort));
    assertThat(appliedNetwork.successors(secondPort), containsInAnyOrder(newOutB));
    assertThat(appliedNetwork.successors(newOutB), containsInAnyOrder(runnerParDoNode));
    assertThat(appliedNetwork.successors(runnerParDoNode), containsInAnyOrder(runnerParDoNodeOut));
    assertThat(
        appliedNetwork.edgesConnecting(firstPort, sdkPortion),
        everyItem(Matchers.<Edges.Edge>instanceOf(HappensBeforeEdge.class)));
    assertThat(
        appliedNetwork.edgesConnecting(sdkPortion, secondPort),
        everyItem(Matchers.<Edges.Edge>instanceOf(HappensBeforeEdge.class)));

    MutableNetwork<Node, Edge> sdkSubnetwork = networkCapture.getValue();
    assertNetworkMaintainsBipartiteStructure(sdkSubnetwork);

    Node sdkNewOutA = Iterables.getOnlyElement(sdkSubnetwork.successors(firstPort));
    Node sdkNewOutB = Iterables.getOnlyElement(sdkSubnetwork.predecessors(secondPort));

    // firstPort -sdkNewOutA-> sdkParDoNode -sdkNewOutB-> secondPort
    assertThat(
        sdkSubnetwork.nodes(),
        containsInAnyOrder(firstPort, sdkNewOutA, sdkParDoNode, sdkNewOutB, secondPort));
    assertThat(sdkSubnetwork.successors(firstPort), containsInAnyOrder(sdkNewOutA));
    assertThat(sdkSubnetwork.successors(sdkNewOutA), containsInAnyOrder(sdkParDoNode));
    assertThat(sdkSubnetwork.successors(sdkParDoNode), containsInAnyOrder(sdkNewOutB));
    assertThat(sdkSubnetwork.successors(sdkNewOutB), containsInAnyOrder(secondPort));
  }

  @Test
  public void testSdkToRunnerToSdkGraph() {
    Node firstSdkPortion = TestNode.create("FirstSdkPortion");
    Node secondSdkPortion = TestNode.create("SecondSdkPortion");
    @SuppressWarnings({"unchecked", "rawtypes"})
    ArgumentCaptor<MutableNetwork<Node, Edge>> networkCapture =
        ArgumentCaptor.forClass((Class) MutableNetwork.class);
    when(registerFnOperationFunction.apply(networkCapture.capture()))
        .thenReturn(firstSdkPortion, secondSdkPortion);

    Node firstPort = TestNode.create("FirstPort");
    Node secondPort = TestNode.create("SecondPort");
    when(portSupplier.apply(anyString(), anyString())).thenReturn(firstPort, secondPort);

    Node readNode = createReadNode("Read", Nodes.ExecutionLocation.SDK_HARNESS);
    Edge readNodeEdge = DefaultEdge.create();
    Node readNodeOut = createInstructionOutputNode("Read.out");
    Edge readNodeOutEdge = DefaultEdge.create();
    Node runnerParDoNode = createParDoNode("RunnerParDo", Nodes.ExecutionLocation.RUNNER_HARNESS);
    Edge runnerParDoNodeEdge = DefaultEdge.create();
    Node runnerParDoNodeOut = createInstructionOutputNode("RunnerParDo.out");
    Edge runnerParDoNodeOutEdge = DefaultEdge.create();
    Node sdkParDoNode = createParDoNode("SdkParDo", Nodes.ExecutionLocation.SDK_HARNESS);
    Edge sdkParDoNodeEdge = DefaultEdge.create();
    Node sdkParDoNodeOut = createInstructionOutputNode("SdkParDo.out");

    // Read -out-> RunnerParDo -out-> SdkParDo
    MutableNetwork<Node, Edge> network = createEmptyNetwork();
    network.addNode(readNode);
    network.addNode(readNodeOut);
    network.addNode(runnerParDoNode);
    network.addNode(runnerParDoNodeOut);
    network.addNode(sdkParDoNodeOut);
    network.addNode(sdkParDoNodeOut);
    network.addEdge(readNode, readNodeOut, readNodeEdge);
    network.addEdge(readNodeOut, runnerParDoNode, readNodeOutEdge);
    network.addEdge(runnerParDoNode, runnerParDoNodeOut, runnerParDoNodeEdge);
    network.addEdge(runnerParDoNodeOut, sdkParDoNode, runnerParDoNodeOutEdge);
    network.addEdge(sdkParDoNode, sdkParDoNodeOut, sdkParDoNodeEdge);

    MutableNetwork<Node, Edge> appliedNetwork =
        createRegisterFnOperation.apply(Graphs.copyOf(network));
    assertNetworkMaintainsBipartiteStructure(appliedNetwork);

    // On each rewire between runner and SDK, we use a new output node
    Node newOutA = Iterables.getOnlyElement(appliedNetwork.successors(firstPort));
    Node newOutB = Iterables.getOnlyElement(appliedNetwork.predecessors(secondPort));

    // firstSdkPortion -> firstPort -newOutA-> RunnerParDo -newOutB-> secondPort -> secondSdkPortion
    assertThat(
        appliedNetwork.nodes(),
        containsInAnyOrder(
            firstSdkPortion,
            firstPort,
            newOutA,
            runnerParDoNode,
            newOutB,
            secondPort,
            secondSdkPortion));
    assertThat(appliedNetwork.successors(firstSdkPortion), containsInAnyOrder(firstPort));
    assertThat(appliedNetwork.successors(firstPort), containsInAnyOrder(newOutA));
    assertThat(appliedNetwork.successors(newOutA), containsInAnyOrder(runnerParDoNode));
    assertThat(appliedNetwork.successors(runnerParDoNode), containsInAnyOrder(newOutB));
    assertThat(appliedNetwork.successors(newOutB), containsInAnyOrder(secondPort));
    assertThat(appliedNetwork.successors(secondPort), containsInAnyOrder(secondSdkPortion));
    assertThat(
        appliedNetwork.edgesConnecting(firstSdkPortion, firstPort),
        everyItem(Matchers.<Edges.Edge>instanceOf(HappensBeforeEdge.class)));
    assertThat(
        appliedNetwork.edgesConnecting(secondPort, secondSdkPortion),
        everyItem(Matchers.<Edges.Edge>instanceOf(HappensBeforeEdge.class)));

    // The order of the calls to create the SDK subnetworks is indeterminate
    List<MutableNetwork<Node, Edge>> sdkSubnetworks = networkCapture.getAllValues();
    MutableNetwork<Node, Edge> firstSdkSubnetwork;
    MutableNetwork<Node, Edge> secondSdkSubnetwork;
    if (sdkSubnetworks.get(0).nodes().contains(readNode)) {
      firstSdkSubnetwork = sdkSubnetworks.get(0);
      secondSdkSubnetwork = sdkSubnetworks.get(1);
    } else {
      firstSdkSubnetwork = sdkSubnetworks.get(1);
      secondSdkSubnetwork = sdkSubnetworks.get(0);
    }
    assertNetworkMaintainsBipartiteStructure(firstSdkSubnetwork);
    assertNetworkMaintainsBipartiteStructure(secondSdkSubnetwork);

    Node sdkNewOutA = Iterables.getOnlyElement(firstSdkSubnetwork.predecessors(firstPort));

    // readNode -sdkNewOutA-> firstPort
    assertThat(firstSdkSubnetwork.nodes(), containsInAnyOrder(readNode, sdkNewOutA, firstPort));
    assertThat(firstSdkSubnetwork.successors(readNode), containsInAnyOrder(sdkNewOutA));
    assertThat(firstSdkSubnetwork.successors(sdkNewOutA), containsInAnyOrder(firstPort));

    Node sdkNewOutB = Iterables.getOnlyElement(secondSdkSubnetwork.successors(secondPort));

    // secondPort -sdkNewOutB-> sdkParDoNode -> sdkParDoNodeOut
    assertThat(
        secondSdkSubnetwork.nodes(),
        containsInAnyOrder(secondPort, sdkNewOutB, sdkParDoNode, sdkParDoNodeOut));
    assertThat(secondSdkSubnetwork.successors(secondPort), containsInAnyOrder(sdkNewOutB));
    assertThat(secondSdkSubnetwork.successors(sdkNewOutB), containsInAnyOrder(sdkParDoNode));
    assertThat(secondSdkSubnetwork.successors(sdkParDoNode), containsInAnyOrder(sdkParDoNodeOut));
  }

  @Test
  public void testRunnerAndSdkToRunnerAndSdkGraph() {
    // RunnerSource --\   /--> RunnerParDo
    //                 out
    // CustomSource --/   \--> SdkParDo
    //
    // Should produce:
    //        PortB --> out --\
    // RunnerSource --> out --> RunnerParDo
    //                     \--> PortA
    //        PortA --> out --\
    // CustomSource --> out --> SdkParDo
    //                     \--> PortB
    Node firstSdkPortion = TestNode.create("FirstSdkPortion");
    Node secondSdkPortion = TestNode.create("SecondSdkPortion");
    @SuppressWarnings({"unchecked", "rawtypes"})
    ArgumentCaptor<MutableNetwork<Node, Edge>> networkCapture =
        ArgumentCaptor.forClass((Class) MutableNetwork.class);
    when(registerFnOperationFunction.apply(networkCapture.capture()))
        .thenReturn(firstSdkPortion, secondSdkPortion);

    Node firstPort = TestNode.create("FirstPort");
    Node secondPort = TestNode.create("SecondPort");
    when(portSupplier.apply(anyString(), anyString())).thenReturn(firstPort, secondPort);

    Node runnerReadNode = createReadNode("RunnerRead", Nodes.ExecutionLocation.RUNNER_HARNESS);
    Edge runnerReadNodeEdge = DefaultEdge.create();
    Node sdkReadNode = createReadNode("SdkRead", Nodes.ExecutionLocation.SDK_HARNESS);
    Edge sdkReadNodeEdge = DefaultEdge.create();
    Node readNodeOut = createInstructionOutputNode("Read.out");
    Edge readNodeOutToRunnerEdge = DefaultEdge.create();
    Edge readNodeOutToSdkEdge = DefaultEdge.create();
    Node runnerParDoNode = createParDoNode("RunnerParDo", Nodes.ExecutionLocation.RUNNER_HARNESS);
    Edge runnerParDoNodeEdge = DefaultEdge.create();
    Node runnerParDoNodeOut = createInstructionOutputNode("RunnerParDo.out");
    Node sdkParDoNode = createParDoNode("SdkParDo", Nodes.ExecutionLocation.SDK_HARNESS);
    Edge sdkParDoNodeEdge = DefaultEdge.create();
    Node sdkParDoNodeOut = createInstructionOutputNode("SdkParDo.out");

    // Read -out-> RunnerParDo -out-> SdkParDo
    MutableNetwork<Node, Edge> network = createEmptyNetwork();
    network.addNode(sdkReadNode);
    network.addNode(runnerReadNode);
    network.addNode(readNodeOut);
    network.addNode(runnerParDoNode);
    network.addNode(runnerParDoNodeOut);
    network.addNode(sdkParDoNodeOut);
    network.addNode(sdkParDoNodeOut);
    network.addEdge(sdkReadNode, readNodeOut, sdkReadNodeEdge);
    network.addEdge(runnerReadNode, readNodeOut, runnerReadNodeEdge);
    network.addEdge(readNodeOut, runnerParDoNode, readNodeOutToRunnerEdge);
    network.addEdge(readNodeOut, sdkParDoNode, readNodeOutToSdkEdge);
    network.addEdge(runnerParDoNode, runnerParDoNodeOut, runnerParDoNodeEdge);
    network.addEdge(sdkParDoNode, sdkParDoNodeOut, sdkParDoNodeEdge);

    MutableNetwork<Node, Edge> appliedNetwork =
        createRegisterFnOperation.apply(Graphs.copyOf(network));
    assertNetworkMaintainsBipartiteStructure(appliedNetwork);

    // Node wiring is indeterministic, must be detected from generated graph.
    Node sdkPortionA;
    Node sdkPortionB;
    if (appliedNetwork.inDegree(firstSdkPortion) == 0) {
      sdkPortionA = firstSdkPortion;
      sdkPortionB = secondSdkPortion;
    } else {
      sdkPortionA = secondSdkPortion;
      sdkPortionB = firstSdkPortion;
    }
    Node portA = Iterables.getOnlyElement(appliedNetwork.successors(sdkPortionA));
    Node portB = Iterables.getOnlyElement(appliedNetwork.predecessors(sdkPortionB));
    // On each rewire between runner and SDK, we use a new output node
    Node newOutA = Iterables.getOnlyElement(appliedNetwork.successors(portA));
    Node newOutB = Iterables.getOnlyElement(appliedNetwork.predecessors(portB));

    // sdkPortionA -> portA -newOutA-> runnerParDoNode -> runnerParDoNodeOut
    //       runnerReadNode -newOutB-/
    //                               \--> portB -> sdkPortionB
    assertThat(
        appliedNetwork.nodes(),
        containsInAnyOrder(
            runnerReadNode,
            firstSdkPortion,
            secondSdkPortion,
            portA,
            newOutA,
            portB,
            newOutB,
            runnerParDoNode,
            runnerParDoNodeOut));
    assertThat(appliedNetwork.successors(runnerReadNode), containsInAnyOrder(newOutB));
    assertThat(appliedNetwork.successors(newOutB), containsInAnyOrder(runnerParDoNode, portB));
    assertThat(appliedNetwork.successors(portB), containsInAnyOrder(sdkPortionB));
    assertThat(appliedNetwork.successors(sdkPortionA), containsInAnyOrder(portA));
    assertThat(appliedNetwork.successors(portA), containsInAnyOrder(newOutA));
    assertThat(appliedNetwork.successors(newOutA), containsInAnyOrder(runnerParDoNode));
    assertThat(appliedNetwork.successors(runnerParDoNode), containsInAnyOrder(runnerParDoNodeOut));
    assertThat(
        appliedNetwork.edgesConnecting(sdkPortionA, portA),
        everyItem(Matchers.<Edges.Edge>instanceOf(HappensBeforeEdge.class)));
    assertThat(
        appliedNetwork.edgesConnecting(portB, sdkPortionB),
        everyItem(Matchers.<Edges.Edge>instanceOf(HappensBeforeEdge.class)));

    // Argument captor call order can be indeterministic
    List<MutableNetwork<Node, Edge>> sdkSubnetworks = networkCapture.getAllValues();
    MutableNetwork<Node, Edge> sdkSubnetworkA;
    MutableNetwork<Node, Edge> sdkSubnetworkB;
    if (sdkSubnetworks.get(0).nodes().contains(sdkReadNode)) {
      sdkSubnetworkA = sdkSubnetworks.get(0);
      sdkSubnetworkB = sdkSubnetworks.get(1);
    } else {
      sdkSubnetworkA = sdkSubnetworks.get(1);
      sdkSubnetworkB = sdkSubnetworks.get(0);
    }
    assertNetworkMaintainsBipartiteStructure(sdkSubnetworkA);
    assertNetworkMaintainsBipartiteStructure(sdkSubnetworkB);

    //                       /-> portA
    // sdkReadNode -sdkNewOutA-> sdkParDoNode -> sdkParDoNodeOut
    Node sdkNewOutA = Iterables.getOnlyElement(sdkSubnetworkA.predecessors(portA));
    assertThat(
        sdkSubnetworkA.nodes(),
        containsInAnyOrder(sdkReadNode, portA, sdkNewOutA, sdkParDoNode, sdkParDoNodeOut));
    assertThat(sdkSubnetworkA.successors(sdkReadNode), containsInAnyOrder(sdkNewOutA));
    assertThat(sdkSubnetworkA.successors(sdkNewOutA), containsInAnyOrder(portA, sdkParDoNode));
    assertThat(sdkSubnetworkA.successors(sdkParDoNode), containsInAnyOrder(sdkParDoNodeOut));

    // portB -sdkNewOutB-> sdkParDoNode -> sdkParDoNodeOut
    Node sdkNewOutB = Iterables.getOnlyElement(sdkSubnetworkB.successors(portB));
    assertThat(
        sdkSubnetworkB.nodes(),
        containsInAnyOrder(portB, sdkNewOutB, sdkParDoNode, sdkParDoNodeOut));
    assertThat(sdkSubnetworkB.successors(portB), containsInAnyOrder(sdkNewOutB));
    assertThat(sdkSubnetworkB.successors(sdkNewOutB), containsInAnyOrder(sdkParDoNode));
    assertThat(sdkSubnetworkB.successors(sdkParDoNode), containsInAnyOrder(sdkParDoNodeOut));
  }

  /**
   * This asserts that all root nodes are not InstructionOutputNodes and that there is a bipartite
   * relationship between non-InstructionOutputNodes and InstructionOutputNodes, ignoring
   * HappensBeforeEdges.
   */
  private static void assertNetworkMaintainsBipartiteStructure(
      Network<Node, Edge> originalNetwork) {
    // Copy the network before we remove happens before relationships.
    MutableNetwork<Node, Edge> network = Graphs.copyOf(originalNetwork);
    for (HappensBeforeEdge edge :
        ImmutableList.copyOf(Iterables.filter(network.edges(), HappensBeforeEdge.class))) {
      network.removeEdge(edge);
    }

    for (Node node : network.nodes()) {
      if (node instanceof InstructionOutputNode) {
        assertTrue(network.inDegree(node) > 0);
        for (Node adjacentNode : network.adjacentNodes(node)) {
          assertThat(adjacentNode, not(instanceOf(InstructionOutputNode.class)));
        }
      } else {
        for (Node adjacentNode : network.adjacentNodes(node)) {
          assertThat(adjacentNode, instanceOf(InstructionOutputNode.class));
        }
      }
    }
  }

  private static MutableNetwork<Node, Edge> createEmptyNetwork() {
    return NetworkBuilder.directed()
        .allowsSelfLoops(false)
        .allowsParallelEdges(true)
        .<Node, Edge>build();
  }

  private static ParallelInstructionNode createReadNode(
      String name, Nodes.ExecutionLocation location) {
    return ParallelInstructionNode.create(
        new ParallelInstruction().setName(name).setRead(new ReadInstruction()), location);
  }

  private static ParallelInstructionNode createParDoNode(
      String name, Nodes.ExecutionLocation location) {
    return ParallelInstructionNode.create(
        new ParallelInstruction().setName(name).setParDo(new ParDoInstruction()), location);
  }

  private static InstructionOutputNode createInstructionOutputNode(String name) {
    return InstructionOutputNode.create(new InstructionOutput().setName(name), "fakeId");
  }

  /** A named node to easily differentiate graph construction problems during testing. */
  @AutoValue
  public abstract static class TestNode extends Node {
    public static TestNode create(String value) {
      return new AutoValue_CreateRegisterFnOperationFunctionTest_TestNode(value);
    }

    public abstract String getName();

    @Override
    public String toString() {
      return hashCode() + " " + getName();
    }
  }
}