@Test
  public void testStorageException() {
    IScheduledTask task = makeTask("a", PENDING);
    TaskInfo mesosTask =
        TaskInfo.newBuilder()
            .setName(Tasks.id(task))
            .setTaskId(TaskID.newBuilder().setValue(Tasks.id(task)))
            .setSlaveId(SlaveID.newBuilder().setValue("slaveId"))
            .build();

    Capture<Runnable> timeoutCapture = expectTaskGroupBackoff(10);
    expectAnyMaintenanceCalls();
    expectOfferDeclineIn(10);
    expect(assigner.maybeAssign(OFFER_A, task)).andThrow(new StorageException("Injected failure."));

    Capture<Runnable> timeoutCapture2 = expectTaskGroupBackoff(10, 20);
    expect(preemptor.findPreemptionSlotFor("a")).andReturn(Optional.<String>absent());
    expect(assigner.maybeAssign(OFFER_A, task)).andReturn(Optional.of(mesosTask));
    driver.launchTask(OFFER_A.getId(), mesosTask);
    expectLastCall();

    replayAndCreateScheduler();

    changeState(task, INIT, PENDING);
    offerQueue.addOffer(OFFER_A);
    timeoutCapture.getValue().run();
    timeoutCapture2.getValue().run();
  }
  @Test
  public void testDriverNotReady() {
    IScheduledTask task = makeTask("a", PENDING);
    TaskInfo mesosTask =
        TaskInfo.newBuilder()
            .setName(Tasks.id(task))
            .setTaskId(TaskID.newBuilder().setValue(Tasks.id(task)))
            .setSlaveId(SlaveID.newBuilder().setValue("slaveId"))
            .build();

    Capture<Runnable> timeoutCapture = expectTaskGroupBackoff(10);
    expectAnyMaintenanceCalls();
    expectOfferDeclineIn(10);
    expect(assigner.maybeAssign(OFFER_A, task)).andReturn(Optional.of(mesosTask));
    driver.launchTask(OFFER_A.getId(), mesosTask);
    expectLastCall().andThrow(new IllegalStateException("Driver not ready."));
    expect(
            stateManager.changeState(
                Query.taskScoped("a").byStatus(PENDING), LOST, TaskScheduler.LAUNCH_FAILED_MSG))
        .andReturn(1);

    replayAndCreateScheduler();

    changeState(task, INIT, PENDING);
    offerQueue.addOffer(OFFER_A);
    timeoutCapture.getValue().run();
  }
  @Test
  public void testTaskAssigned() {
    expectAnyMaintenanceCalls();
    expectOfferDeclineIn(10);

    IScheduledTask task = makeTask("a", PENDING);
    TaskInfo mesosTask = makeTaskInfo(task);

    Capture<Runnable> timeoutCapture = expectTaskGroupBackoff(10);
    expect(assigner.maybeAssign(OFFER_A, task)).andReturn(Optional.<TaskInfo>absent());
    expect(preemptor.findPreemptionSlotFor("a")).andReturn(Optional.<String>absent());

    Capture<Runnable> timeoutCapture2 = expectTaskGroupBackoff(10, 20);
    expect(assigner.maybeAssign(OFFER_A, task)).andReturn(Optional.of(mesosTask));
    driver.launchTask(OFFER_A.getId(), mesosTask);

    Capture<Runnable> timeoutCapture3 = expectTaskGroupBackoff(10);
    expectTaskGroupBackoff(10, 20);
    expect(preemptor.findPreemptionSlotFor("b")).andReturn(Optional.<String>absent());

    replayAndCreateScheduler();

    offerQueue.addOffer(OFFER_A);
    changeState(task, INIT, PENDING);
    timeoutCapture.getValue().run();
    timeoutCapture2.getValue().run();

    // Ensure the offer was consumed.
    changeState(makeTask("b"), INIT, PENDING);
    timeoutCapture3.getValue().run();
  }
 private Capture<IScheduledTask> expectTaskScheduled(IScheduledTask task) {
   TaskInfo mesosTask = makeTaskInfo(task);
   Capture<IScheduledTask> taskScheduled = createCapture();
   expect(assigner.maybeAssign(EasyMock.<Offer>anyObject(), capture(taskScheduled)))
       .andReturn(Optional.of(mesosTask));
   driver.launchTask(EasyMock.<OfferID>anyObject(), eq(mesosTask));
   return taskScheduled;
 }
  @Test
  public void testChangingMaintenancePreferences() {
    expectOffer();
    expect(maintenance.getMode("HOST_A")).andReturn(MaintenanceMode.NONE);
    expectOffer();
    expect(maintenance.getMode("HOST_B")).andReturn(MaintenanceMode.SCHEDULED);
    expectOffer();
    expect(maintenance.getMode("HOST_C")).andReturn(MaintenanceMode.DRAINED);

    IScheduledTask taskA = makeTask("A", PENDING);
    TaskInfo mesosTaskA = makeTaskInfo(taskA);
    expect(assigner.maybeAssign(OFFER_B, taskA)).andReturn(Optional.of(mesosTaskA));
    driver.launchTask(OFFER_B.getId(), mesosTaskA);
    Capture<Runnable> captureA = expectTaskGroupBackoff(10);

    IScheduledTask taskB = makeTask("B", PENDING);
    TaskInfo mesosTaskB = makeTaskInfo(taskB);
    expect(assigner.maybeAssign(OFFER_C, taskB)).andReturn(Optional.of(mesosTaskB));
    driver.launchTask(OFFER_C.getId(), mesosTaskB);
    Capture<Runnable> captureB = expectTaskGroupBackoff(10);

    replayAndCreateScheduler();

    offerQueue.addOffer(OFFER_A);
    offerQueue.addOffer(OFFER_B);
    offerQueue.addOffer(OFFER_C);

    // Initially, we'd expect the offers to be consumed in order (A, B), with (C) unschedulable

    // Expected order now (B), with (C, A) unschedulable
    changeHostMaintenanceState("HOST_A", MaintenanceMode.DRAINING);
    changeState(taskA, INIT, PENDING);
    captureA.getValue().run();

    // Expected order now (C), with (A) unschedulable and (B) already consumed
    changeHostMaintenanceState("HOST_C", MaintenanceMode.NONE);
    changeState(taskB, INIT, PENDING);
    captureB.getValue().run();
  }
  @Test
  public void testBasicMaintenancePreferences() {
    expectOffer();
    expect(maintenance.getMode("HOST_D")).andReturn(MaintenanceMode.DRAINED);
    expectOffer();
    expect(maintenance.getMode("HOST_C")).andReturn(MaintenanceMode.DRAINING);
    expectOffer();
    expect(maintenance.getMode("HOST_B")).andReturn(MaintenanceMode.SCHEDULED);
    expectOffer();
    expect(maintenance.getMode("HOST_A")).andReturn(MaintenanceMode.NONE);

    IScheduledTask taskA = makeTask("A", PENDING);
    TaskInfo mesosTaskA = makeTaskInfo(taskA);
    expect(assigner.maybeAssign(OFFER_A, taskA)).andReturn(Optional.of(mesosTaskA));
    driver.launchTask(OFFER_A.getId(), mesosTaskA);
    Capture<Runnable> captureA = expectTaskGroupBackoff(10);

    IScheduledTask taskB = makeTask("B", PENDING);
    TaskInfo mesosTaskB = makeTaskInfo(taskB);
    expect(assigner.maybeAssign(OFFER_B, taskB)).andReturn(Optional.of(mesosTaskB));
    driver.launchTask(OFFER_B.getId(), mesosTaskB);
    Capture<Runnable> captureB = expectTaskGroupBackoff(10);

    replayAndCreateScheduler();

    offerQueue.addOffer(OFFER_D);
    offerQueue.addOffer(OFFER_C);
    offerQueue.addOffer(OFFER_B);
    offerQueue.addOffer(OFFER_A);

    changeState(taskA, INIT, PENDING);
    captureA.getValue().run();

    changeState(taskB, INIT, PENDING);
    captureB.getValue().run();
  }
  @Test
  public void testDontDeclineAcceptedOffer() throws OfferQueue.LaunchException {
    expectAnyMaintenanceCalls();
    Capture<Runnable> offerExpirationCapture = expectOfferDeclineIn(10);

    Function<Offer, Optional<TaskInfo>> offerAcceptor =
        createMock(new Clazz<Function<Offer, Optional<TaskInfo>>>() {});
    final TaskInfo taskInfo = TaskInfo.getDefaultInstance();
    expect(offerAcceptor.apply(OFFER_A)).andReturn(Optional.of(taskInfo));
    driver.launchTask(OFFER_A.getId(), taskInfo);

    replayAndCreateScheduler();

    offerQueue.addOffer(OFFER_A);
    offerQueue.launchFirst(offerAcceptor);
    offerExpirationCapture.getValue().run();
  }