@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();
}