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 testRMContainerAllocatorExceptionIsHandled() throws Exception { ClientService mockClientService = mock(ClientService.class); AppContext mockContext = mock(AppContext.class); MockRMCommunicator mockRMCommunicator = new MockRMCommunicator(mockClientService, mockContext); RMCommunicator communicator = spy(mockRMCommunicator); Clock mockClock = mock(Clock.class); when(mockContext.getClock()).thenReturn(mockClock); doThrow(new RMContainerAllocationException("Test")).doNothing() .when(communicator).heartbeat(); when(mockClock.getTime()).thenReturn(1L).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(); }
@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( (Answer<Long>) invocation -> { communicator.stop(); return 2L; }).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(); }