Java Code Examples for org.apache.tez.dag.api.EdgeProperty.SchedulingType#SEQUENTIAL

@Test(expected = IllegalStateException.class, timeout = 5000)
public void testInputAndInputVertexNameCollision() {
  Vertex v1 = new Vertex("v1",
      new ProcessorDescriptor("MapProcessor"),
      dummyTaskCount, dummyTaskResource);
  Vertex v2 = new Vertex("v2",
      new ProcessorDescriptor("MapProcessor"),
      dummyTaskCount, dummyTaskResource);
  
  v2.addInput("v1", new InputDescriptor(), null);
  
  Edge e1 = new Edge(v1, v2,
      new EdgeProperty(DataMovementType.SCATTER_GATHER, 
          DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, 
          new OutputDescriptor("dummy output class"),
          new InputDescriptor("dummy input class")));
  
  DAG dag = new DAG("testDag");
  dag.addVertex(v1);
  dag.addVertex(v2);
  dag.addEdge(e1);
  dag.verify();
}
 

@Test(expected = IllegalStateException.class, timeout = 5000)
public void testOutputAndOutputVertexNameCollision() {
  Vertex v1 = new Vertex("v1",
      new ProcessorDescriptor("MapProcessor"),
      dummyTaskCount, dummyTaskResource);
  Vertex v2 = new Vertex("v2",
      new ProcessorDescriptor("MapProcessor"),
      dummyTaskCount, dummyTaskResource);
  
  v1.addOutput("v2", new OutputDescriptor());
  
  Edge e1 = new Edge(v1, v2,
      new EdgeProperty(DataMovementType.SCATTER_GATHER, 
          DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, 
          new OutputDescriptor("dummy output class"),
          new InputDescriptor("dummy input class")));
  
  DAG dag = new DAG("testDag");
  dag.addVertex(v1);
  dag.addVertex(v2);
  dag.addEdge(e1);
  dag.verify();
}
 

@Test(expected = IllegalStateException.class, timeout = 5000)  
public void testVerify3() {
  Vertex v1 = new Vertex("v1",
      new ProcessorDescriptor(dummyProcessorClassName),
      dummyTaskCount, dummyTaskResource);
  Vertex v2 = new Vertex("v2",
      new ProcessorDescriptor("MapProcessor"),
      dummyTaskCount, dummyTaskResource);
  Edge e1 = new Edge(v1, v2,
      new EdgeProperty(DataMovementType.SCATTER_GATHER, 
          DataSourceType.EPHEMERAL, SchedulingType.SEQUENTIAL, 
          new OutputDescriptor(dummyOutputClassName),
          new InputDescriptor(dummyInputClassName)));
  DAG dag = new DAG("testDag");
  dag.addVertex(v1);
  dag.addVertex(v2);
  dag.addEdge(e1);
  dag.verify();
}
 

@Test(timeout = 5000)
public void testVerifyBroadcast() {
  Vertex v1 = new Vertex("v1",
      new ProcessorDescriptor(dummyProcessorClassName),
      dummyTaskCount, dummyTaskResource);
  Vertex v2 = new Vertex("v2",
      new ProcessorDescriptor("MapProcessor"),
      dummyTaskCount, dummyTaskResource);
  Edge e1 = new Edge(v1, v2,
      new EdgeProperty(DataMovementType.BROADCAST, 
          DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, 
          new OutputDescriptor(dummyOutputClassName),
          new InputDescriptor(dummyInputClassName)));
  DAG dag = new DAG("testDag");
  dag.addVertex(v1);
  dag.addVertex(v2);
  dag.addEdge(e1);
  dag.verify();
}
 

@Test(timeout = 5000)
public void testVerifyOneToOneNoInferParallelism() {
  Vertex v1 = new Vertex("v1",
      new ProcessorDescriptor(dummyProcessorClassName),
      -1, dummyTaskResource);
  Vertex v2 = new Vertex("v2",
      new ProcessorDescriptor("MapProcessor"),
      -1, dummyTaskResource);
  Edge e1 = new Edge(v1, v2,
      new EdgeProperty(DataMovementType.ONE_TO_ONE, 
          DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, 
          new OutputDescriptor(dummyOutputClassName),
          new InputDescriptor(dummyInputClassName)));
  DAG dag = new DAG("testDag");
  dag.addVertex(v1);
  dag.addVertex(v2);
  dag.addEdge(e1);
  dag.verify();
  Assert.assertEquals(-1, v2.getParallelism());
}
 

@Test(timeout = 5000)
public void testVerifyScatterGather() {
  Vertex v1 = new Vertex("v1",
      new ProcessorDescriptor(dummyProcessorClassName),
      dummyTaskCount, dummyTaskResource);
  Vertex v2 = new Vertex("v2",
      new ProcessorDescriptor("MapProcessor"),
      dummyTaskCount, dummyTaskResource);
  Edge e1 = new Edge(v1, v2,
      new EdgeProperty(DataMovementType.SCATTER_GATHER, 
          DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, 
          new OutputDescriptor(dummyOutputClassName),
          new InputDescriptor(dummyInputClassName)));
  DAG dag = new DAG("testDag");
  dag.addVertex(v1);
  dag.addVertex(v2);
  dag.addEdge(e1);
  dag.verify();
}
 

@Test(timeout = 5000)
public void BinaryOutput() {
  IllegalStateException ex = null;
  try {
    Vertex v1 = new Vertex("v1",
        new ProcessorDescriptor("MapProcessor"),
        dummyTaskCount, dummyTaskResource);
    Vertex v2 = new Vertex("v2",
        new ProcessorDescriptor("MapProcessor"),
        dummyTaskCount, dummyTaskResource);
    Vertex v3 = new Vertex("v3",
        new ProcessorDescriptor("MapProcessor"),
        dummyTaskCount, dummyTaskResource);
    Edge e1 = new Edge(v1, v2,
        new EdgeProperty(DataMovementType.SCATTER_GATHER, 
            DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, 
            new OutputDescriptor("dummy output class"),
            new InputDescriptor("dummy input class")));
    Edge e2 = new Edge(v1, v2,
        new EdgeProperty(DataMovementType.SCATTER_GATHER, 
            DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, 
            new OutputDescriptor("dummy output class"),
            new InputDescriptor("dummy input class")));
    DAG dag = new DAG("testDag");
    dag.addVertex(v1);
    dag.addVertex(v2);
    dag.addVertex(v3);
    dag.addEdge(e1);
    dag.addEdge(e2);
    dag.verify();
  }
  catch (IllegalStateException e){
    ex = e;
  }
  Assert.assertNull(ex);
}
 

public TezEdgeDescriptor() {
    combinePlan = new PhysicalPlan();

    // The default is shuffle edge.
    inputClassName = OrderedGroupedKVInput.class.getName();
    outputClassName = OrderedPartitionedKVOutput.class.getName();
    partitionerClass = null;
    schedulingType = SchedulingType.SEQUENTIAL;
    dataSourceType = DataSourceType.PERSISTED;
    dataMovementType = DataMovementType.SCATTER_GATHER;
}
 

public static SchedulingType convertFromDAGPlan(PlanEdgeSchedulingType type){
  switch(type){
    case SEQUENTIAL : return SchedulingType.SEQUENTIAL;
    case CONCURRENT : return SchedulingType.CONCURRENT;
    default : throw new IllegalArgumentException("unknown 'SchedulingType': " + type);
  }
}
 

@Test(timeout = 5000)
// v1 (known) -> v2 (-1) -> v3 (-1)
public void testVerifyOneToOneInferParallelism() {
  Vertex v1 = new Vertex("v1",
      new ProcessorDescriptor(dummyProcessorClassName),
      dummyTaskCount, dummyTaskResource);
  Vertex v2 = new Vertex("v2",
      new ProcessorDescriptor("MapProcessor"),
      -1, dummyTaskResource);
  Vertex v3 = new Vertex("v3",
      new ProcessorDescriptor("MapProcessor"),
      -1, dummyTaskResource);
  Edge e1 = new Edge(v1, v2,
      new EdgeProperty(DataMovementType.ONE_TO_ONE, 
          DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, 
          new OutputDescriptor(dummyOutputClassName),
          new InputDescriptor(dummyInputClassName)));
  Edge e2 = new Edge(v2, v3,
      new EdgeProperty(DataMovementType.ONE_TO_ONE, 
          DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, 
          new OutputDescriptor(dummyOutputClassName),
          new InputDescriptor(dummyInputClassName)));
  DAG dag = new DAG("testDag");
  dag.addVertex(v1);
  dag.addVertex(v2);
  dag.addVertex(v3);
  dag.addEdge(e1);
  dag.addEdge(e2);
  dag.verify();
  Assert.assertEquals(dummyTaskCount, v2.getParallelism());
  Assert.assertEquals(dummyTaskCount, v3.getParallelism());
}
 

@Test(timeout = 5000)
public void BinaryInputAllowed() {
  Vertex v1 = new Vertex("v1",
      new ProcessorDescriptor("MapProcessor"),
      dummyTaskCount, dummyTaskResource);
  Vertex v2 = new Vertex("v2",
      new ProcessorDescriptor("MapProcessor"),
      dummyTaskCount, dummyTaskResource);
  Vertex v3 = new Vertex("v3",
      new ProcessorDescriptor("ReduceProcessor"),
      dummyTaskCount, dummyTaskResource);
  Edge e1 = new Edge(v1, v3,
      new EdgeProperty(DataMovementType.SCATTER_GATHER, 
          DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL,
          new OutputDescriptor("dummy output class"),
          new InputDescriptor("dummy input class")));
  Edge e2 = new Edge(v2, v3,
      new EdgeProperty(DataMovementType.SCATTER_GATHER, 
          DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, 
          new OutputDescriptor("dummy output class"),
          new InputDescriptor("dummy input class")));
  DAG dag = new DAG("testDag");
  dag.addVertex(v1);
  dag.addVertex(v2);
  dag.addVertex(v3);
  dag.addEdge(e1);
  dag.addEdge(e2);
  dag.verify();
}
 

public static SchedulingType convertFromDAGPlan(PlanEdgeSchedulingType type){
  switch(type){
    case SEQUENTIAL : return SchedulingType.SEQUENTIAL;
    case CONCURRENT : return SchedulingType.CONCURRENT;
    default : throw new IllegalArgumentException("unknown 'SchedulingType': " + type);
  }
}
 

@Test(timeout = 5000)
// v1 (-1) -> v2 (known) -> v3 (-1)
public void testVerifyOneToOneIncorrectParallelism1() {
  Vertex v1 = new Vertex("v1",
      new ProcessorDescriptor(dummyProcessorClassName),
      -1, dummyTaskResource);
  Vertex v2 = new Vertex("v2",
      new ProcessorDescriptor(dummyProcessorClassName),
      dummyTaskCount, dummyTaskResource);
  Vertex v3 = new Vertex("v3",
      new ProcessorDescriptor("MapProcessor"),
      -1, dummyTaskResource);
  Edge e1 = new Edge(v1, v3,
      new EdgeProperty(DataMovementType.ONE_TO_ONE, 
          DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, 
          new OutputDescriptor(dummyOutputClassName),
          new InputDescriptor(dummyInputClassName)));
  Edge e2 = new Edge(v2, v3,
      new EdgeProperty(DataMovementType.ONE_TO_ONE, 
          DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, 
          new OutputDescriptor(dummyOutputClassName),
          new InputDescriptor(dummyInputClassName)));
  DAG dag = new DAG("testDag");
  dag.addVertex(v1);
  dag.addVertex(v2);
  dag.addVertex(v3);
  dag.addEdge(e1);
  dag.addEdge(e2);
  try {
    dag.verify();
    Assert.assertTrue(false);
  } catch (TezUncheckedException e) {
    Assert.assertTrue(e.getMessage().contains(
        "1-1 Edge. Destination vertex parallelism must match source vertex"));
  }
}
 

@Test(timeout = 5000)
public void testUserPayloadSerde() {
  DAG dag = new DAG("testDag");
  ProcessorDescriptor pd1 = new ProcessorDescriptor("processor1").
      setUserPayload("processor1Bytes".getBytes());
  ProcessorDescriptor pd2 = new ProcessorDescriptor("processor2").
      setUserPayload("processor2Bytes".getBytes());
  Vertex v1 = new Vertex("v1", pd1, 10, Resource.newInstance(1024, 1));
  Vertex v2 = new Vertex("v2", pd2, 1, Resource.newInstance(1024, 1));
  v1.setTaskLaunchCmdOpts("").setTaskEnvironment(new HashMap<String, String>())
      .setTaskLocalFiles(new HashMap<String, LocalResource>());
  v2.setTaskLaunchCmdOpts("").setTaskEnvironment(new HashMap<String, String>())
      .setTaskLocalFiles(new HashMap<String, LocalResource>());

  InputDescriptor inputDescriptor = new InputDescriptor("input").
      setUserPayload("inputBytes".getBytes());
  OutputDescriptor outputDescriptor = new OutputDescriptor("output").
      setUserPayload("outputBytes".getBytes());
  Edge edge = new Edge(v1, v2, new EdgeProperty(
      DataMovementType.SCATTER_GATHER, DataSourceType.PERSISTED,
      SchedulingType.SEQUENTIAL, outputDescriptor, inputDescriptor));

  dag.addVertex(v1).addVertex(v2).addEdge(edge);

  DAGPlan dagProto = dag.createDag(new TezConfiguration());

  assertEquals(2, dagProto.getVertexCount());
  assertEquals(1, dagProto.getEdgeCount());

  VertexPlan v1Proto = dagProto.getVertex(0);
  VertexPlan v2Proto = dagProto.getVertex(1);
  EdgePlan edgeProto = dagProto.getEdge(0);

  assertEquals("processor1Bytes", new String(v1Proto.getProcessorDescriptor()
      .getUserPayload().toByteArray()));
  assertEquals("processor1", v1Proto.getProcessorDescriptor().getClassName());

  assertEquals("processor2Bytes", new String(v2Proto.getProcessorDescriptor()
      .getUserPayload().toByteArray()));
  assertEquals("processor2", v2Proto.getProcessorDescriptor().getClassName());

  assertEquals("inputBytes", new String(edgeProto.getEdgeDestination()
      .getUserPayload().toByteArray()));
  assertEquals("input", edgeProto.getEdgeDestination().getClassName());

  assertEquals("outputBytes", new String(edgeProto.getEdgeSource()
      .getUserPayload().toByteArray()));
  assertEquals("output", edgeProto.getEdgeSource().getClassName());

  EdgeProperty edgeProperty = DagTypeConverters
      .createEdgePropertyMapFromDAGPlan(dagProto.getEdgeList().get(0));

  byte[] ib = edgeProperty.getEdgeDestination().getUserPayload();
  assertEquals("inputBytes", new String(ib));
  assertEquals("input", edgeProperty.getEdgeDestination().getClassName());

  byte[] ob = edgeProperty.getEdgeSource().getUserPayload();
  assertEquals("outputBytes", new String(ob));
  assertEquals("output", edgeProperty.getEdgeSource().getClassName());
}
 

@Test (timeout=5000)
public void testCredentialsSerde() {
  DAG dag = new DAG("testDag");
  ProcessorDescriptor pd1 = new ProcessorDescriptor("processor1").
      setUserPayload("processor1Bytes".getBytes());
  ProcessorDescriptor pd2 = new ProcessorDescriptor("processor2").
      setUserPayload("processor2Bytes".getBytes());
  Vertex v1 = new Vertex("v1", pd1, 10, Resource.newInstance(1024, 1));
  Vertex v2 = new Vertex("v2", pd2, 1, Resource.newInstance(1024, 1));
  v1.setTaskLaunchCmdOpts("").setTaskEnvironment(new HashMap<String, String>())
      .setTaskLocalFiles(new HashMap<String, LocalResource>());
  v2.setTaskLaunchCmdOpts("").setTaskEnvironment(new HashMap<String, String>())
      .setTaskLocalFiles(new HashMap<String, LocalResource>());

  InputDescriptor inputDescriptor = new InputDescriptor("input").
      setUserPayload("inputBytes".getBytes());
  OutputDescriptor outputDescriptor = new OutputDescriptor("output").
      setUserPayload("outputBytes".getBytes());
  Edge edge = new Edge(v1, v2, new EdgeProperty(
      DataMovementType.SCATTER_GATHER, DataSourceType.PERSISTED,
      SchedulingType.SEQUENTIAL, outputDescriptor, inputDescriptor));

  dag.addVertex(v1).addVertex(v2).addEdge(edge);

  Credentials dagCredentials = new Credentials();
  Token<TokenIdentifier> token1 = new Token<TokenIdentifier>();
  Token<TokenIdentifier> token2 = new Token<TokenIdentifier>();
  dagCredentials.addToken(new Text("Token1"), token1);
  dagCredentials.addToken(new Text("Token2"), token2);
  
  dag.setCredentials(dagCredentials);

  DAGPlan dagProto = dag.createDag(new TezConfiguration());

  assertTrue(dagProto.hasCredentialsBinary());
  
  Credentials fetchedCredentials = DagTypeConverters.convertByteStringToCredentials(dagProto
      .getCredentialsBinary());
  
  assertEquals(2, fetchedCredentials.numberOfTokens());
  assertNotNull(fetchedCredentials.getToken(new Text("Token1")));
  assertNotNull(fetchedCredentials.getToken(new Text("Token2")));
}
 

@Test(timeout = 5000)
public void testVertexGroup() {
  Vertex v1 = new Vertex("v1",
      new ProcessorDescriptor("Processor"),
      dummyTaskCount, dummyTaskResource);
  Vertex v2 = new Vertex("v2",
      new ProcessorDescriptor("Processor"),
      dummyTaskCount, dummyTaskResource);
  Vertex v3 = new Vertex("v3",
      new ProcessorDescriptor("Processor"),
      dummyTaskCount, dummyTaskResource);
  Vertex v4 = new Vertex("v4",
      new ProcessorDescriptor("Processor"),
      dummyTaskCount, dummyTaskResource);
  Vertex v5 = new Vertex("v5",
      new ProcessorDescriptor("Processor"),
      dummyTaskCount, dummyTaskResource);
  
  DAG dag = new DAG("testDag");
  String groupName1 = "uv12";
  VertexGroup uv12 = dag.createVertexGroup(groupName1, v1, v2);
  OutputDescriptor outDesc = new OutputDescriptor();
  uv12.addOutput("uvOut", outDesc, null);
  
  String groupName2 = "uv23";
  VertexGroup uv23 = dag.createVertexGroup(groupName2, v2, v3);
  
  GroupInputEdge e1 = new GroupInputEdge(uv12, v4,
      new EdgeProperty(DataMovementType.SCATTER_GATHER, 
          DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL,
          new OutputDescriptor("dummy output class"),
          new InputDescriptor("dummy input class")),
          new InputDescriptor("dummy input class"));
  GroupInputEdge e2 = new GroupInputEdge(uv23, v5,
      new EdgeProperty(DataMovementType.SCATTER_GATHER, 
          DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL,
          new OutputDescriptor("dummy output class"),
          new InputDescriptor("dummy input class")),
          new InputDescriptor("dummy input class"));
  
  dag.addVertex(v1);
  dag.addVertex(v2);
  dag.addVertex(v3);
  dag.addVertex(v4);
  dag.addVertex(v5);
  dag.addEdge(e1);
  dag.addEdge(e2);
  dag.verify();

  // for the first Group v1 and v2 should get connected to v4 and also have 1 output
  // for the second Group v2 and v3 should get connected to v5
  // the Group place holders should disappear
  Assert.assertNull(dag.getVertex(uv12.getGroupName()));
  Assert.assertNull(dag.getVertex(uv23.getGroupName()));
  Assert.assertFalse(dag.edges.contains(e1));
  Assert.assertFalse(dag.edges.contains(e2));
  Assert.assertEquals(1, v1.getOutputs().size());
  Assert.assertEquals(1, v2.getOutputs().size());
  Assert.assertEquals(outDesc, v1.getOutputs().get(0).getDescriptor());
  Assert.assertEquals(outDesc, v2.getOutputs().get(0).getDescriptor());
  Assert.assertEquals(1, v1.getOutputVertices().size());
  Assert.assertEquals(1, v3.getOutputVertices().size());
  Assert.assertEquals(2, v2.getOutputVertices().size());
  Assert.assertTrue(v1.getOutputVertices().contains(v4));
  Assert.assertTrue(v3.getOutputVertices().contains(v5));
  Assert.assertTrue(v2.getOutputVertices().contains(v4));
  Assert.assertTrue(v2.getOutputVertices().contains(v5));
  Assert.assertEquals(2, v4.getInputVertices().size());
  Assert.assertTrue(v4.getInputVertices().contains(v1));
  Assert.assertTrue(v4.getInputVertices().contains(v2));
  Assert.assertEquals(2, v5.getInputVertices().size());
  Assert.assertTrue(v5.getInputVertices().contains(v2));
  Assert.assertTrue(v5.getInputVertices().contains(v3));
  Assert.assertEquals(1, v4.getGroupInputs().size());
  Assert.assertTrue(v4.getGroupInputs().containsKey(groupName1));
  Assert.assertEquals(1, v5.getGroupInputs().size());
  Assert.assertTrue(v5.getGroupInputs().containsKey(groupName2));
  Assert.assertEquals(2, dag.vertexGroups.size());
}
 

@Test(timeout = 5000)
public void testVertexGroupOneToOne() {
  Vertex v1 = new Vertex("v1",
      new ProcessorDescriptor("Processor"),
      dummyTaskCount, dummyTaskResource);
  Vertex v2 = new Vertex("v2",
      new ProcessorDescriptor("Processor"),
      dummyTaskCount, dummyTaskResource);
  Vertex v3 = new Vertex("v3",
      new ProcessorDescriptor("Processor"),
      dummyTaskCount, dummyTaskResource);
  Vertex v4 = new Vertex("v4",
      new ProcessorDescriptor("Processor"),
      dummyTaskCount, dummyTaskResource);
  Vertex v5 = new Vertex("v5",
      new ProcessorDescriptor("Processor"),
      -1, dummyTaskResource);
  
  DAG dag = new DAG("testDag");
  String groupName1 = "uv12";
  VertexGroup uv12 = dag.createVertexGroup(groupName1, v1, v2);
  OutputDescriptor outDesc = new OutputDescriptor();
  uv12.addOutput("uvOut", outDesc, null);
  
  String groupName2 = "uv23";
  VertexGroup uv23 = dag.createVertexGroup(groupName2, v2, v3);
  
  GroupInputEdge e1 = new GroupInputEdge(uv12, v4,
      new EdgeProperty(DataMovementType.ONE_TO_ONE, 
          DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL,
          new OutputDescriptor("dummy output class"),
          new InputDescriptor("dummy input class")),
          new InputDescriptor("dummy input class"));
  GroupInputEdge e2 = new GroupInputEdge(uv23, v5,
      new EdgeProperty(DataMovementType.ONE_TO_ONE, 
          DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL,
          new OutputDescriptor("dummy output class"),
          new InputDescriptor("dummy input class")),
          new InputDescriptor("dummy input class"));
  
  dag.addVertex(v1);
  dag.addVertex(v2);
  dag.addVertex(v3);
  dag.addVertex(v4);
  dag.addVertex(v5);
  dag.addEdge(e1);
  dag.addEdge(e2);
  dag.verify();
  Assert.assertEquals(dummyTaskCount, v5.getParallelism());
}
 

@SuppressWarnings({ "rawtypes", "unchecked" })
@Test (timeout = 5000)
public void testCompositeEventHandling() {
  EventHandler eventHandler = mock(EventHandler.class);
  EdgeProperty edgeProp = new EdgeProperty(DataMovementType.SCATTER_GATHER,
      DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, mock(OutputDescriptor.class),
      mock(InputDescriptor.class));
  Edge edge = new Edge(edgeProp, eventHandler);
  
  TezVertexID srcVertexID = createVertexID(1);
  TezVertexID destVertexID = createVertexID(2);
  LinkedHashMap<TezTaskID, Task> srcTasks = mockTasks(srcVertexID, 1);
  LinkedHashMap<TezTaskID, Task> destTasks = mockTasks(destVertexID, 5);
  
  TezTaskID srcTaskID = srcTasks.keySet().iterator().next();
  
  Vertex srcVertex = mockVertex("src", srcVertexID, srcTasks);
  Vertex destVertex = mockVertex("dest", destVertexID, destTasks);
  
  edge.setSourceVertex(srcVertex);
  edge.setDestinationVertex(destVertex);
  edge.initialize();
  
  TezTaskAttemptID srcTAID = createTAIDForTest(srcTaskID, 2); // Task0, Attempt 0
  
  EventMetaData srcMeta = new EventMetaData(EventProducerConsumerType.OUTPUT, "consumerVertex", "producerVertex", srcTAID);
  
  // Verification via a CompositeEvent
  CompositeDataMovementEvent cdmEvent = new CompositeDataMovementEvent(0, destTasks.size(), "bytes".getBytes());
  cdmEvent.setVersion(2); // AttemptNum
  TezEvent tezEvent = new TezEvent(cdmEvent, srcMeta);
  // Event setup to look like it would after the Vertex is done with it.

  edge.sendTezEventToDestinationTasks(tezEvent);
  
  ArgumentCaptor<Event> args = ArgumentCaptor.forClass(Event.class);
  verify(eventHandler, times(destTasks.size())).handle(args.capture());
  
  verifyEvents(args.getAllValues(), srcTAID, destTasks);
  
  
  // Same Verification via regular DataMovementEvents
  reset(eventHandler);
  for (int i = 0 ; i < destTasks.size() ; i++) {
    DataMovementEvent dmEvent = new DataMovementEvent(i, "bytes".getBytes());
    dmEvent.setVersion(2);
    tezEvent = new TezEvent(dmEvent, srcMeta);
    edge.sendTezEventToDestinationTasks(tezEvent);
  }
  args = ArgumentCaptor.forClass(Event.class);
  verify(eventHandler, times(destTasks.size())).handle(args.capture());
  
  verifyEvents(args.getAllValues(), srcTAID, destTasks);

}
 

@Test (timeout=5000)
public void testBasicOneToOne() {
  HashMap<String, EdgeProperty> mockInputVertices = 
      new HashMap<String, EdgeProperty>();
  String mockSrcVertexId1 = "Vertex1";
  EdgeProperty eProp1 = new EdgeProperty(
      EdgeProperty.DataMovementType.ONE_TO_ONE,
      EdgeProperty.DataSourceType.PERSISTED,
      SchedulingType.SEQUENTIAL, 
      new OutputDescriptor("out"),
      new InputDescriptor("in"));
  
  String mockManagedVertexId = "Vertex";
  Container mockContainer1 = mock(Container.class);
  ContainerId mockCId1 = mock(ContainerId.class);
  when(mockContainer1.getId()).thenReturn(mockCId1);
  Container mockContainer2 = mock(Container.class);
  ContainerId mockCId2 = mock(ContainerId.class);
  when(mockContainer2.getId()).thenReturn(mockCId2);
  
  VertexManagerPluginContext mockContext = mock(VertexManagerPluginContext.class);
  when(mockContext.getInputVertexEdgeProperties()).thenReturn(mockInputVertices);
  when(mockContext.getVertexName()).thenReturn(mockManagedVertexId);
  when(mockContext.getVertexNumTasks(mockManagedVertexId)).thenReturn(3);
  when(mockContext.getVertexNumTasks(mockSrcVertexId1)).thenReturn(3);
  when(mockContext.getTaskContainer(mockSrcVertexId1, 0)).thenReturn(mockContainer1);
  when(mockContext.getTaskContainer(mockSrcVertexId1, 1)).thenReturn(mockContainer2);
  when(mockContext.getTaskContainer(mockSrcVertexId1, 2)).thenReturn(mockContainer1);
  mockInputVertices.put(mockSrcVertexId1, eProp1);
  
  Map<String, List<Integer>> initialCompletions = Maps.newHashMap();
  initialCompletions.put(mockSrcVertexId1, Collections.singletonList(0));
  
  InputReadyVertexManager manager = new InputReadyVertexManager();
  manager.initialize(mockContext);
  manager.onVertexStarted(initialCompletions);
  verify(mockContext, times(1)).scheduleVertexTasks(requestCaptor.capture());
  Assert.assertEquals(1, requestCaptor.getValue().size());
  Assert.assertEquals(0, requestCaptor.getValue().get(0).getTaskIndex().intValue());
  Assert.assertEquals(mockCId1, requestCaptor.getValue().get(0)
      .getTaskLocationHint().getAffinitizedContainer());
  manager.onSourceTaskCompleted(mockSrcVertexId1, 1);
  verify(mockContext, times(2)).scheduleVertexTasks(requestCaptor.capture());
  Assert.assertEquals(1, requestCaptor.getValue().size());
  Assert.assertEquals(1, requestCaptor.getValue().get(0).getTaskIndex().intValue());
  Assert.assertEquals(mockCId2, requestCaptor.getValue().get(0)
      .getTaskLocationHint().getAffinitizedContainer());
  manager.onSourceTaskCompleted(mockSrcVertexId1, 2);
  verify(mockContext, times(3)).scheduleVertexTasks(requestCaptor.capture());
  Assert.assertEquals(1, requestCaptor.getValue().size());
  Assert.assertEquals(2, requestCaptor.getValue().get(0).getTaskIndex().intValue());
  Assert.assertEquals(mockCId1, requestCaptor.getValue().get(0)
      .getTaskLocationHint().getAffinitizedContainer());
}
 

@Test(timeout = 5000)
public void testVertexGroupWithMultipleOutputEdges() {
  Vertex v1 = new Vertex("v1",
      new ProcessorDescriptor("Processor"),
      dummyTaskCount, dummyTaskResource);
  Vertex v2 = new Vertex("v2",
      new ProcessorDescriptor("Processor"),
      dummyTaskCount, dummyTaskResource);
  Vertex v3 = new Vertex("v3",
      new ProcessorDescriptor("Processor"),
      dummyTaskCount, dummyTaskResource);
  Vertex v4 = new Vertex("v4",
      new ProcessorDescriptor("Processor"),
      dummyTaskCount, dummyTaskResource);
  
  DAG dag = new DAG("testDag");
  VertexGroup uv12 = dag.createVertexGroup("uv12", v1, v2);
  OutputDescriptor outDesc = new OutputDescriptor();
  uv12.addOutput("uvOut", outDesc, null);
  
  GroupInputEdge e1 = new GroupInputEdge(uv12, v3,
      new EdgeProperty(DataMovementType.SCATTER_GATHER, 
          DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL,
          new OutputDescriptor("dummy output class"),
          new InputDescriptor("dummy input class")),
          new InputDescriptor("dummy input class"));
  
  GroupInputEdge e2 = new GroupInputEdge(uv12, v4,
      new EdgeProperty(DataMovementType.SCATTER_GATHER, 
          DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL,
          new OutputDescriptor("dummy output class"),
          new InputDescriptor("dummy input class")),
          new InputDescriptor("dummy input class"));

  dag.addVertex(v1);
  dag.addVertex(v2);
  dag.addVertex(v3);
  dag.addVertex(v4);
  dag.addEdge(e1);
  dag.addEdge(e2);
  dag.verify();
  
  Assert.assertEquals(2, v1.getOutputVertices().size());
  Assert.assertEquals(2, v2.getOutputVertices().size());
  Assert.assertTrue(v1.getOutputVertices().contains(v3));
  Assert.assertTrue(v1.getOutputVertices().contains(v4));
  Assert.assertTrue(v2.getOutputVertices().contains(v3));
  Assert.assertTrue(v2.getOutputVertices().contains(v4));
}