TaskStatusHandlerImpl extends AbstractExecutionThreadService implements TaskStatusHandler { @Override public void statusUpdate(TaskStatus status) { pendingUpdates.add(status); } @Inject  TaskStatusHandlerImpl(
      Storage storage,
      StateManager stateManager,
      final Driver driver,
      @StatusUpdateQueue BlockingQueue<TaskStatus> pendingUpdates,
      @MaxBatchSize Integer maxBatchSize,
      CachedCounters counters); @Override void statusUpdate(TaskStatus status);  }
@Test public void testFailedStatusUpdate() throws Exception { storageUtil.expectWrite(); CountDownLatch latch = new CountDownLatch(1); expect(stateManager.changeState( storageUtil.mutableStoreProvider, TASK_ID_A, Optional.absent(), RUNNING, Optional.of("fake message"))) .andAnswer(() -> { latch.countDown(); throw new StorageException("Injected error"); }); control.replay(); TaskStatus status = TaskStatus.newBuilder() .setState(TaskState.TASK_RUNNING) .setTaskId(TaskID.newBuilder().setValue(TASK_ID_A)) .setMessage("fake message") .build(); statusHandler.statusUpdate(status); assertTrue(latch.await(5L, TimeUnit.SECONDS)); }
@Test public void testSuppressUnregisteredExecutorMessage() throws Exception { storageUtil.expectWrite(); TaskStatus status = TaskStatus.newBuilder() .setState(TaskState.TASK_KILLED) .setTaskId(TaskID.newBuilder().setValue(TASK_ID_A)) .setReason(TaskStatus.Reason.REASON_EXECUTOR_UNREGISTERED) .setMessage("Unregistered executor") .build(); expect(stateManager.changeState( storageUtil.mutableStoreProvider, TASK_ID_A, Optional.absent(), KILLED, Optional.absent())) .andReturn(StateChangeResult.SUCCESS); CountDownLatch latch = new CountDownLatch(1); driver.acknowledgeStatusUpdate(status); waitAndAnswer(latch); control.replay(); statusHandler.statusUpdate(status); assertTrue(latch.await(5L, TimeUnit.SECONDS)); }
@Test public void testThreadFailure() throws Exception { statusHandler.stopAsync(); statusHandler.awaitTerminated(); stateManager = createMock(StateManager.class); storageUtil = new StorageTestUtil(this); driver = createMock(Driver.class); queue = createMock(new Clazz<BlockingQueue<TaskStatus>>() { }); statusHandler = new TaskStatusHandlerImpl( storageUtil.storage, stateManager, driver, queue, 1000, new CachedCounters(stats)); expect(queue.add(EasyMock.anyObject())).andReturn(true); expect(queue.take()).andAnswer(() -> { throw new RuntimeException(); }); CountDownLatch latch = new CountDownLatch(1); driver.abort(); waitAndAnswer(latch); control.replay(); statusHandler.startAsync(); TaskStatus status = TaskStatus.newBuilder() .setState(TaskState.TASK_RUNNING) .setTaskId(TaskID.newBuilder().setValue(TASK_ID_A)) .setMessage("fake message") .build(); statusHandler.statusUpdate(status); assertTrue(latch.await(5L, TimeUnit.SECONDS)); }