StateManagerImpl implements StateManager { @Override public StateChangeResult changeState( MutableStoreProvider storeProvider, String taskId, Optional<ScheduleStatus> casState, final ScheduleStatus newState, final Optional<String> auditMessage) { return updateTaskAndExternalState( storeProvider.getUnsafeTaskStore(), casState, taskId, newState, auditMessage); } @Inject  StateManagerImpl(
      final Clock clock,
      Driver driver,
      TaskIdGenerator taskIdGenerator,
      EventSink eventSink,
      RescheduleCalculator rescheduleCalculator); @Override void insertPendingTasks(
      MutableStoreProvider storeProvider,
      final ITaskConfig task,
      Set<Integer> instanceIds); @Override StateChangeResult changeState(
      MutableStoreProvider storeProvider,
      String taskId,
      Optional<ScheduleStatus> casState,
      final ScheduleStatus newState,
      final Optional<String> auditMessage); @Override IAssignedTask assignTask(
      MutableStoreProvider storeProvider,
      String taskId,
      String slaveHost,
      SlaveID slaveId,
      Map<String, Integer> assignedPorts); @Override void deleteTasks(MutableStoreProvider storeProvider, final Set<String> taskIds);  }
@Test public void testKillPendingTask() { String taskId = "a"; expect(taskIdGenerator.generate(NON_SERVICE_CONFIG, 0)).andReturn(taskId); expectStateTransitions(taskId, INIT, PENDING); eventSink.post(matchTasksDeleted(taskId)); control.replay(); insertTask(NON_SERVICE_CONFIG, 0); assertEquals(SUCCESS, changeState(taskId, KILLING)); assertEquals(ILLEGAL, changeState(taskId, KILLING)); }
@Test public void testNestedEvents() { final String id = "a"; expect(taskIdGenerator.generate(NON_SERVICE_CONFIG, 0)).andReturn(id); eventSink.post(matchStateChange(id, INIT, PENDING)); expectLastCall().andAnswer(() -> { changeState(id, ASSIGNED); return null; }); expectStateTransitions(id, PENDING, ASSIGNED); control.replay(); insertTask(NON_SERVICE_CONFIG, 0); }
@Test public void testDeletePendingTask() { String taskId = "a"; expect(taskIdGenerator.generate(NON_SERVICE_CONFIG, 0)).andReturn(taskId); expectStateTransitions(taskId, INIT, PENDING); eventSink.post(matchTasksDeleted(taskId)); control.replay(); insertTask(NON_SERVICE_CONFIG, 0); changeState(taskId, KILLING); }
@Test public void testThrottleTask() { ITaskConfig task = setIsService(NON_SERVICE_CONFIG, true); String taskId = "a"; expect(taskIdGenerator.generate(task, 0)).andReturn(taskId); expectStateTransitions(taskId, INIT, PENDING, ASSIGNED, RUNNING, FAILED); String newTaskId = "b"; expect(taskIdGenerator.generate(task, 0)).andReturn(newTaskId); expect(rescheduleCalculator.getFlappingPenaltyMs(EasyMock.anyObject())).andReturn(100L); expectStateTransitions(newTaskId, INIT, THROTTLED); control.replay(); insertTask(task, 0); changeState(taskId, ASSIGNED); changeState(taskId, RUNNING); changeState(taskId, FAILED); }
@Test public void testKillUnknownTask() { String unknownTask = "unknown"; driver.killTask(unknownTask); control.replay(); changeState(unknownTask, RUNNING); }
@Test public void testCasTaskNotFound() { control.replay(); assertEquals(INVALID_CAS_STATE, changeState( "a", Optional.of(PENDING), ASSIGNED, Optional.absent())); }