RMCommunicator extends AbstractService implements RMHeartbeatHandler { protected abstract void heartbeat() throws Exception; RMCommunicator(ClientService clientService, AppContext context); @Override long getLastHeartbeatTime(); @Override void runOnNextHeartbeat(Runnable callback); void setShouldUnregister(boolean shouldUnregister); void setSignalled(boolean isSignalled); EnumSet<SchedulerResourceTypes> getSchedulerResourceTypes();  }
@Test(timeout = 2000) public void testRMContainerAllocatorYarnRuntimeExceptionIsHandled() throws Exception { ClientService mockClientService = mock(ClientService.class); AppContext mockContext = mock(AppContext.class); MockRMCommunicator mockRMCommunicator = new MockRMCommunicator(mockClientService, mockContext); final RMCommunicator communicator = spy(mockRMCommunicator); Clock mockClock = mock(Clock.class); when(mockContext.getClock()).thenReturn(mockClock); doThrow(new YarnRuntimeException("Test")).doNothing() .when(communicator).heartbeat(); when(mockClock.getTime()).thenReturn(1L).thenAnswer(new Answer<Integer>() { @Override public Integer answer(InvocationOnMock invocation) throws Throwable { communicator.stop(); return 2; } }).thenThrow(new AssertionError( "GetClock called second time, when it should not have since the thread " + "should have quit")); AllocatorRunnable testRunnable = communicator.new AllocatorRunnable(); testRunnable.run(); verify(mockClock, times(2)).getTime(); }