@Test(timeout = 10000)
public void testEventThrottling() throws Exception {
TezTaskAttemptID mockTaskAttemptId = mock(TezTaskAttemptID.class);
LogicalIOProcessorRuntimeTask mockTask = mock(LogicalIOProcessorRuntimeTask.class);
when(mockTask.getMaxEventsToHandle()).thenReturn(10000, 1);
when(mockTask.getVertexName()).thenReturn("vertexName");
when(mockTask.getTaskAttemptID()).thenReturn(mockTaskAttemptId);
TezTaskUmbilicalProtocol mockUmbilical = mock(TezTaskUmbilicalProtocol.class);
TezHeartbeatResponse resp1 = new TezHeartbeatResponse(createEvents(5));
resp1.setLastRequestId(1);
TezHeartbeatResponse resp2 = new TezHeartbeatResponse(createEvents(1));
resp2.setLastRequestId(2);
resp2.setShouldDie();
when(mockUmbilical.heartbeat(isA(TezHeartbeatRequest.class))).thenReturn(resp1, resp2);
// Setup the sleep time to be way higher than the test timeout
TaskReporter.HeartbeatCallable heartbeatCallable =
new TaskReporter.HeartbeatCallable(
mockTask, mockUmbilical, 100000, 100000, 5, new AtomicLong(0), "containerIdStr");
ExecutorService executor = Executors.newSingleThreadExecutor();
try {
Future<Boolean> result = executor.submit(heartbeatCallable);
Assert.assertFalse(result.get());
} finally {
executor.shutdownNow();
}
ArgumentCaptor<TezHeartbeatRequest> captor = ArgumentCaptor.forClass(TezHeartbeatRequest.class);
verify(mockUmbilical, times(2)).heartbeat(captor.capture());
TezHeartbeatRequest req = captor.getValue();
Assert.assertEquals(2, req.getRequestId());
Assert.assertEquals(1, req.getMaxEvents());
}
