StateManagerImpl implements StateManager { @Override public IAssignedTask assignTask( MutableStoreProvider storeProvider, String taskId, String slaveHost, SlaveID slaveId, Map<String, Integer> assignedPorts) { checkNotBlank(taskId); checkNotBlank(slaveHost); requireNonNull(slaveId); requireNonNull(assignedPorts); IScheduledTask mutated = storeProvider.getUnsafeTaskStore().mutateTask(taskId, task -> { ScheduledTask builder = task.newBuilder(); builder.getAssignedTask() .setAssignedPorts(assignedPorts) .setSlaveHost(slaveHost) .setSlaveId(slaveId.getValue()); return IScheduledTask.build(builder); }).get(); StateChangeResult changeResult = updateTaskAndExternalState( storeProvider.getUnsafeTaskStore(), Optional.absent(), taskId, ASSIGNED, Optional.absent()); Preconditions.checkState( changeResult == SUCCESS, "Attempt to assign task " + taskId + " to " + slaveHost + " failed"); return mutated.getAssignedTask(); } @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 testPortResource() throws Exception { Set<String> requestedPorts = ImmutableSet.of("one", "two", "three"); ITaskConfig task = setRequestedPorts(NON_SERVICE_CONFIG, requestedPorts); String taskId = "a"; expect(taskIdGenerator.generate(task, 0)).andReturn(taskId); expectStateTransitions(taskId, INIT, PENDING, ASSIGNED); control.replay(); insertTask(task, 0); assignTask(taskId, HOST_A, ImmutableMap.of("one", 80, "two", 81, "three", 82)); IScheduledTask actual = Storage.Util.fetchTask(storage, taskId).get(); assertEquals( requestedPorts, actual.getAssignedTask().getTask().getRequestedPorts()); }
@Test public void testAssignTaskPubsub() { String taskId = "a"; expect(taskIdGenerator.generate(NON_SERVICE_CONFIG, 0)).andReturn(taskId); expectStateTransitions(taskId, INIT, PENDING); Capture<TaskStateChange> taskStateChangeCapture = createCapture(); eventSink.post(capture(taskStateChangeCapture)); control.replay(); insertTask(NON_SERVICE_CONFIG, 0); assignTask(taskId, HOST_A); TaskStateChange change = taskStateChangeCapture.getValue(); assertEquals(ASSIGNED, change.getNewState()); assertEquals(HOST_A.getHost(), change.getTask().getAssignedTask().getSlaveHost()); assertEquals(HOST_A.getSlaveId(), change.getTask().getAssignedTask().getSlaveId()); }