ShuffleVertexManager scheduler = null;
EventHandler mockEventHandler = mock(EventHandler.class);
TezDAGID dagId = new TezDAGID("1", 1, 1);
HashMap<Vertex, Edge> mockInputVertices =
new HashMap<Vertex, Edge>();
Vertex mockSrcVertex1 = mock(Vertex.class);
TezVertexID mockSrcVertexId1 = new TezVertexID(dagId, 1);
EdgeProperty eProp1 = new EdgeProperty(
EdgeProperty.DataMovementType.SCATTER_GATHER,
EdgeProperty.DataSourceType.PERSISTED,
SchedulingType.SEQUENTIAL,
new OutputDescriptor("out"),
new InputDescriptor("in"));
when(mockSrcVertex1.getVertexId()).thenReturn(mockSrcVertexId1);
Vertex mockSrcVertex2 = mock(Vertex.class);
TezVertexID mockSrcVertexId2 = new TezVertexID(dagId, 2);
EdgeProperty eProp2 = new EdgeProperty(
EdgeProperty.DataMovementType.SCATTER_GATHER,
EdgeProperty.DataSourceType.PERSISTED,
SchedulingType.SEQUENTIAL,
new OutputDescriptor("out"),
new InputDescriptor("in"));
when(mockSrcVertex2.getVertexId()).thenReturn(mockSrcVertexId2);
Vertex mockSrcVertex3 = mock(Vertex.class);
TezVertexID mockSrcVertexId3 = new TezVertexID(dagId, 3);
EdgeProperty eProp3 = new EdgeProperty(
EdgeProperty.DataMovementType.BROADCAST,
EdgeProperty.DataSourceType.PERSISTED,
SchedulingType.SEQUENTIAL,
new OutputDescriptor("out"),
new InputDescriptor("in"));
when(mockSrcVertex3.getVertexId()).thenReturn(mockSrcVertexId3);
Vertex mockManagedVertex = mock(Vertex.class);
TezVertexID mockManagedVertexId = new TezVertexID(dagId, 3);
when(mockManagedVertex.getVertexId()).thenReturn(mockManagedVertexId);
when(mockManagedVertex.getInputVertices()).thenReturn(mockInputVertices);
// fail if there is no bipartite src vertex
mockInputVertices.put(mockSrcVertex3, new Edge(eProp3, mockEventHandler));
try {
scheduler = createScheduler(conf, mockManagedVertex, 0.1f, 0.1f);
Assert.assertFalse(true);
} catch (TezUncheckedException e) {
Assert.assertTrue(e.getMessage().contains(
"Atleast 1 bipartite source should exist"));
}
mockInputVertices.put(mockSrcVertex1, new Edge(eProp1, mockEventHandler));
mockInputVertices.put(mockSrcVertex2, new Edge(eProp2, mockEventHandler));
// check initialization
scheduler = createScheduler(conf, mockManagedVertex, 0.1f, 0.1f);
Assert.assertTrue(scheduler.bipartiteSources.size() == 2);
Assert.assertTrue(scheduler.bipartiteSources.containsKey(mockSrcVertexId1));
Assert.assertTrue(scheduler.bipartiteSources.containsKey(mockSrcVertexId2));
HashMap<TezTaskID, Task> managedTasks = new HashMap<TezTaskID, Task>();
TezTaskID mockTaskId1 = new TezTaskID(mockManagedVertexId, 0);
managedTasks.put(mockTaskId1, null);
TezTaskID mockTaskId2 = new TezTaskID(mockManagedVertexId, 1);
managedTasks.put(mockTaskId2, null);
TezTaskID mockTaskId3 = new TezTaskID(mockManagedVertexId, 2);
managedTasks.put(mockTaskId3, null);
when(mockManagedVertex.getTotalTasks()).thenReturn(3);
when(mockManagedVertex.getTasks()).thenReturn(managedTasks);
final HashSet<TezTaskID> scheduledTasks = new HashSet<TezTaskID>();
doAnswer(new Answer() {
public Object answer(InvocationOnMock invocation) {
Object[] args = invocation.getArguments();