// TODO remove hacky name lookup
@Override
public boolean needsWaitAfterOutputConsumable() {
Vertex vertex = getVertex();
ProcessorDescriptor processorDescriptor = vertex.getProcessorDescriptor();
if (processorDescriptor != null
&& processorDescriptor.getClassName().contains("InitialTaskWithInMemSort")) {
return true;
} else {
return false;
}
}
@Before
public void setUp() {
dispatcher = new DrainDispatcher();
dispatcher.register(DAGEventType.class, mock(EventHandler.class));
dispatcher.register(VertexEventType.class, mock(EventHandler.class));
dispatcher.register(TaskEventType.class, new TaskEventHandler());
dispatcher.register(TaskAttemptEventType.class, taEventHandler);
dispatcher.init(new Configuration());
dispatcher.start();
vertex = mock(Vertex.class, RETURNS_DEEP_STUBS);
when(vertex.getProcessorDescriptor().getClassName()).thenReturn("");
mockAppContext = mock(AppContext.class, RETURNS_DEEP_STUBS);
when(mockAppContext.getCurrentDAG().getVertex(any(TezVertexID.class))).thenReturn(vertex);
mockHistoryEventHandler = new MockHistoryEventHandler(mockAppContext);
when(mockAppContext.getHistoryHandler()).thenReturn(mockHistoryEventHandler);
task =
new TaskImpl(
vertexId,
0,
dispatcher.getEventHandler(),
new Configuration(),
mock(TaskCommunicatorManagerInterface.class),
new SystemClock(),
mock(TaskHeartbeatHandler.class),
mockAppContext,
false,
Resource.newInstance(1, 1),
mock(ContainerContext.class),
mock(StateChangeNotifier.class),
vertex);
Map<String, OutputCommitter> committers = new HashMap<String, OutputCommitter>();
committers.put(
"out1", new TestOutputCommitter(mock(OutputCommitterContext.class), true, false));
when(task.getVertex().getOutputCommitters()).thenReturn(committers);
}