/**
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.aurora.scheduler.thrift;

import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Proxy;
import java.util.Map;
import java.util.Optional;
import java.util.Set;

import com.google.common.base.Strings;
import com.google.common.base.Throwables;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Iterables;
import com.google.common.collect.Sets;

import org.apache.aurora.common.testing.easymock.EasyMockTest;
import org.apache.aurora.gen.AssignedTask;
import org.apache.aurora.gen.AuroraAdmin;
import org.apache.aurora.gen.Constraint;
import org.apache.aurora.gen.Container;
import org.apache.aurora.gen.CoordinatorSlaPolicy;
import org.apache.aurora.gen.ExecutorConfig;
import org.apache.aurora.gen.ExplicitReconciliationSettings;
import org.apache.aurora.gen.HostStatus;
import org.apache.aurora.gen.Hosts;
import org.apache.aurora.gen.InstanceTaskConfig;
import org.apache.aurora.gen.JobConfiguration;
import org.apache.aurora.gen.JobUpdate;
import org.apache.aurora.gen.JobUpdateInstructions;
import org.apache.aurora.gen.JobUpdatePulseStatus;
import org.apache.aurora.gen.JobUpdateQuery;
import org.apache.aurora.gen.JobUpdateRequest;
import org.apache.aurora.gen.JobUpdateSettings;
import org.apache.aurora.gen.JobUpdateStrategy;
import org.apache.aurora.gen.JobUpdateSummary;
import org.apache.aurora.gen.LimitConstraint;
import org.apache.aurora.gen.ListBackupsResult;
import org.apache.aurora.gen.MaintenanceMode;
import org.apache.aurora.gen.MesosContainer;
import org.apache.aurora.gen.Metadata;
import org.apache.aurora.gen.PercentageSlaPolicy;
import org.apache.aurora.gen.PulseJobUpdateResult;
import org.apache.aurora.gen.QueryRecoveryResult;
import org.apache.aurora.gen.QueueJobUpdateStrategy;
import org.apache.aurora.gen.Range;
import org.apache.aurora.gen.ReadOnlyScheduler;
import org.apache.aurora.gen.Resource;
import org.apache.aurora.gen.ResourceAggregate;
import org.apache.aurora.gen.Response;
import org.apache.aurora.gen.ResponseCode;
import org.apache.aurora.gen.ResponseDetail;
import org.apache.aurora.gen.Result;
import org.apache.aurora.gen.ScheduleStatus;
import org.apache.aurora.gen.ScheduledTask;
import org.apache.aurora.gen.SlaPolicy;
import org.apache.aurora.gen.StartJobUpdateResult;
import org.apache.aurora.gen.TaskConfig;
import org.apache.aurora.gen.TaskConstraint;
import org.apache.aurora.gen.TaskQuery;
import org.apache.aurora.gen.ValueConstraint;
import org.apache.aurora.gen.VariableBatchJobUpdateStrategy;
import org.apache.aurora.gen.apiConstants;
import org.apache.aurora.scheduler.base.Query;
import org.apache.aurora.scheduler.base.TaskTestUtil;
import org.apache.aurora.scheduler.base.Tasks;
import org.apache.aurora.scheduler.configuration.ConfigurationManager;
import org.apache.aurora.scheduler.configuration.ConfigurationManager.TaskDescriptionException;
import org.apache.aurora.scheduler.configuration.SanitizedConfiguration;
import org.apache.aurora.scheduler.cron.CronException;
import org.apache.aurora.scheduler.cron.CronJobManager;
import org.apache.aurora.scheduler.cron.SanitizedCronJob;
import org.apache.aurora.scheduler.maintenance.MaintenanceController;
import org.apache.aurora.scheduler.quota.QuotaCheckResult;
import org.apache.aurora.scheduler.quota.QuotaManager;
import org.apache.aurora.scheduler.reconciliation.TaskReconciler;
import org.apache.aurora.scheduler.state.StateChangeResult;
import org.apache.aurora.scheduler.state.StateManager;
import org.apache.aurora.scheduler.state.UUIDGenerator;
import org.apache.aurora.scheduler.storage.SnapshotStore;
import org.apache.aurora.scheduler.storage.Storage.StorageException;
import org.apache.aurora.scheduler.storage.backup.Recovery;
import org.apache.aurora.scheduler.storage.backup.StorageBackup;
import org.apache.aurora.scheduler.storage.entities.IHostStatus;
import org.apache.aurora.scheduler.storage.entities.IJobConfiguration;
import org.apache.aurora.scheduler.storage.entities.IJobUpdate;
import org.apache.aurora.scheduler.storage.entities.IMetadata;
import org.apache.aurora.scheduler.storage.entities.IRange;
import org.apache.aurora.scheduler.storage.entities.IResourceAggregate;
import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
import org.apache.aurora.scheduler.storage.entities.ITaskConfig;
import org.apache.aurora.scheduler.storage.testing.StorageTestUtil;
import org.apache.aurora.scheduler.testing.FakeStatsProvider;
import org.apache.aurora.scheduler.updater.JobUpdateController;
import org.apache.aurora.scheduler.updater.JobUpdateController.AuditData;
import org.apache.aurora.scheduler.updater.JobUpdateController.JobUpdatingException;
import org.apache.aurora.scheduler.updater.UpdateInProgressException;
import org.apache.aurora.scheduler.updater.UpdateStateException;
import org.apache.thrift.TException;
import org.easymock.EasyMock;
import org.easymock.IExpectationSetters;
import org.junit.Before;
import org.junit.Test;

import static org.apache.aurora.gen.MaintenanceMode.DRAINING;
import static org.apache.aurora.gen.MaintenanceMode.NONE;
import static org.apache.aurora.gen.MaintenanceMode.SCHEDULED;
import static org.apache.aurora.gen.Resource.diskMb;
import static org.apache.aurora.gen.Resource.numCpus;
import static org.apache.aurora.gen.Resource.ramMb;
import static org.apache.aurora.gen.ResponseCode.INVALID_REQUEST;
import static org.apache.aurora.gen.ResponseCode.JOB_UPDATING_ERROR;
import static org.apache.aurora.gen.ResponseCode.OK;
import static org.apache.aurora.scheduler.configuration.ConfigurationManager.DEDICATED_ATTRIBUTE;
import static org.apache.aurora.scheduler.storage.backup.Recovery.RecoveryException;
import static org.apache.aurora.scheduler.thrift.Fixtures.CRON_JOB;
import static org.apache.aurora.scheduler.thrift.Fixtures.ENOUGH_QUOTA;
import static org.apache.aurora.scheduler.thrift.Fixtures.IDENTITY;
import static org.apache.aurora.scheduler.thrift.Fixtures.INSTANCE_KEY;
import static org.apache.aurora.scheduler.thrift.Fixtures.INVALID_TASK_CONFIG;
import static org.apache.aurora.scheduler.thrift.Fixtures.JOB_KEY;
import static org.apache.aurora.scheduler.thrift.Fixtures.JOB_NAME;
import static org.apache.aurora.scheduler.thrift.Fixtures.NOT_ENOUGH_QUOTA;
import static org.apache.aurora.scheduler.thrift.Fixtures.ROLE;
import static org.apache.aurora.scheduler.thrift.Fixtures.TASK_ID;
import static org.apache.aurora.scheduler.thrift.Fixtures.UPDATE_KEY;
import static org.apache.aurora.scheduler.thrift.Fixtures.USER;
import static org.apache.aurora.scheduler.thrift.Fixtures.UU_ID;
import static org.apache.aurora.scheduler.thrift.Fixtures.assertOkResponse;
import static org.apache.aurora.scheduler.thrift.Fixtures.assertResponse;
import static org.apache.aurora.scheduler.thrift.Fixtures.defaultTask;
import static org.apache.aurora.scheduler.thrift.Fixtures.makeJob;
import static org.apache.aurora.scheduler.thrift.Fixtures.makeProdJob;
import static org.apache.aurora.scheduler.thrift.Fixtures.nonProductionTask;
import static org.apache.aurora.scheduler.thrift.Fixtures.okResponse;
import static org.apache.aurora.scheduler.thrift.Fixtures.productionTask;
import static org.apache.aurora.scheduler.thrift.Fixtures.response;
import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.ADD_INSTANCES;
import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.CREATE_JOB;
import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.CREATE_OR_UPDATE_CRON;
import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.DRAIN_HOSTS;
import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.END_MAINTENANCE;
import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.INVALID_SLA_AWARE_UPDATE;
import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.KILL_TASKS;
import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.MAINTENANCE_STATUS;
import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.NOOP_JOB_UPDATE_MESSAGE;
import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.NO_CRON;
import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.PRUNE_TASKS;
import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.RESTART_SHARDS;
import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.SLA_DRAIN_HOSTS;
import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.START_JOB_UPDATE;
import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.START_MAINTENANCE;
import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.jobAlreadyExistsMessage;
import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.noCronScheduleMessage;
import static org.apache.aurora.scheduler.thrift.SchedulerThriftInterface.notScheduledCronMessage;
import static org.easymock.EasyMock.anyInt;
import static org.easymock.EasyMock.anyObject;
import static org.easymock.EasyMock.eq;
import static org.easymock.EasyMock.expect;
import static org.easymock.EasyMock.expectLastCall;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;

public class SchedulerThriftInterfaceTest extends EasyMockTest {

  private static final String AUDIT_MESSAGE = "message";
  private static final AuditData AUDIT = new AuditData(USER, Optional.of(AUDIT_MESSAGE));
  private static final Thresholds THRESHOLDS = new Thresholds(1000, 2000);
  private static final String EXECUTOR_NAME = apiConstants.AURORA_EXECUTOR_NAME;
  private static final ImmutableSet<Metadata> METADATA =
      ImmutableSet.of(new Metadata("k1", "v1"), new Metadata("k2", "v2"));

  private StorageTestUtil storageUtil;
  private SnapshotStore snapshotStore;
  private StorageBackup backup;
  private Recovery recovery;
  private MaintenanceController maintenance;
  private AuroraAdmin.Iface thrift;
  private CronJobManager cronJobManager;
  private QuotaManager quotaManager;
  private StateManager stateManager;
  private UUIDGenerator uuidGenerator;
  private JobUpdateController jobUpdateController;
  private ReadOnlyScheduler.Iface readOnlyScheduler;
  private AuditMessages auditMessages;
  private TaskReconciler taskReconciler;
  private FakeStatsProvider statsProvider;

  @Before
  public void setUp() throws Exception {
    storageUtil = new StorageTestUtil(this);
    storageUtil.expectOperations();
    snapshotStore = createMock(SnapshotStore.class);
    backup = createMock(StorageBackup.class);
    recovery = createMock(Recovery.class);
    maintenance = createMock(MaintenanceController.class);
    cronJobManager = createMock(CronJobManager.class);
    quotaManager = createMock(QuotaManager.class);
    stateManager = createMock(StateManager.class);
    uuidGenerator = createMock(UUIDGenerator.class);
    jobUpdateController = createMock(JobUpdateController.class);
    readOnlyScheduler = createMock(ReadOnlyScheduler.Iface.class);
    auditMessages = createMock(AuditMessages.class);
    taskReconciler = createMock(TaskReconciler.class);
    statsProvider = new FakeStatsProvider();

    thrift = getResponseProxy(
        new SchedulerThriftInterface(
            TaskTestUtil.CONFIGURATION_MANAGER,
            THRESHOLDS,
            storageUtil.storage,
            snapshotStore,
            backup,
            recovery,
            cronJobManager,
            maintenance,
            quotaManager,
            stateManager,
            uuidGenerator,
            jobUpdateController,
            readOnlyScheduler,
            auditMessages,
            taskReconciler,
            statsProvider));
  }

  private static AuroraAdmin.Iface getResponseProxy(AuroraAdmin.Iface realThrift) {
    // Capture all API method calls to validate response objects.
    Class<AuroraAdmin.Iface> thriftClass = AuroraAdmin.Iface.class;
    return (AuroraAdmin.Iface) Proxy.newProxyInstance(
        thriftClass.getClassLoader(),
        new Class<?>[] {thriftClass},
        (o, method, args) -> {
          Response response;
          try {
            response = (Response) method.invoke(realThrift, args);
          } catch (InvocationTargetException e) {
            Throwables.propagateIfPossible(e.getTargetException(), TException.class);
            throw e;
          }
          assertTrue(response.isSetResponseCode());
          assertNotNull(response.getDetails());
          return response;
        });
  }

  private static SanitizedConfiguration fromUnsanitized(IJobConfiguration job)
      throws TaskDescriptionException {

    return SanitizedConfiguration.fromUnsanitized(TaskTestUtil.CONFIGURATION_MANAGER, job);
  }

  @Test
  public void testCreateJob() throws Exception {
    IJobConfiguration job = IJobConfiguration.build(makeProdJob());
    SanitizedConfiguration sanitized = fromUnsanitized(job);
    jobUpdateController.assertNotUpdating(JOB_KEY);
    storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active());
    expectNoCronJob();
    expectInstanceQuotaCheck(sanitized, ENOUGH_QUOTA);

    stateManager.insertPendingTasks(
        storageUtil.mutableStoreProvider,
        sanitized.getJobConfig().getTaskConfig(),
        sanitized.getInstanceIds());

    control.replay();

    assertOkResponse(thrift.createJob(makeProdJob()));
    assertEquals(1L, statsProvider.getLongValue(CREATE_JOB));
  }

  @Test
  public void testCreateJobFailsForCron() throws Exception {
    IJobConfiguration job = IJobConfiguration.build(makeProdJob().setCronSchedule(""));

    control.replay();

    assertEquals(invalidResponse(NO_CRON), thrift.createJob(job.newBuilder()));
    assertEquals(0L, statsProvider.getLongValue(CREATE_JOB));
  }

  @Test
  public void testCreateJobFailsConfigCheck() throws Exception {
    IJobConfiguration job = IJobConfiguration.build(makeJob(INVALID_TASK_CONFIG));
    control.replay();

    assertResponse(INVALID_REQUEST, thrift.createJob(job.newBuilder()));
    assertEquals(0L, statsProvider.getLongValue(CREATE_JOB));
  }

  @Test
  public void testCreateJobFailsJobExists() throws Exception {
    IJobConfiguration job = IJobConfiguration.build(makeJob());
    jobUpdateController.assertNotUpdating(JOB_KEY);
    storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active(), buildScheduledTask());

    control.replay();

    assertResponse(INVALID_REQUEST, thrift.createJob(job.newBuilder()));
    assertEquals(0L, statsProvider.getLongValue(CREATE_JOB));
  }

  @Test
  public void testCreateJobFailsCronJobExists() throws Exception {
    IJobConfiguration job = IJobConfiguration.build(makeJob());
    jobUpdateController.assertNotUpdating(JOB_KEY);
    storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active());
    expectCronJob();

    control.replay();

    assertResponse(INVALID_REQUEST, thrift.createJob(job.newBuilder()));
    assertEquals(0L, statsProvider.getLongValue(CREATE_JOB));
  }

  @Test
  public void testCreateJobFailsInstanceCheck() throws Exception {
    IJobConfiguration job = IJobConfiguration.build(
        makeJob(defaultTask(true), THRESHOLDS.getMaxTasksPerJob() + 1));

    jobUpdateController.assertNotUpdating(JOB_KEY);
    storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active());
    expectNoCronJob();
    expect(quotaManager.checkInstanceAddition(
        anyObject(ITaskConfig.class),
        anyInt(),
        eq(storageUtil.mutableStoreProvider))).andReturn(ENOUGH_QUOTA);

    control.replay();

    assertResponse(INVALID_REQUEST, thrift.createJob(job.newBuilder()));
    assertEquals(0L, statsProvider.getLongValue(CREATE_JOB));
  }

  @Test
  public void testCreateJobFailsQuotaCheck() throws Exception {
    IJobConfiguration job = IJobConfiguration.build(makeProdJob());
    SanitizedConfiguration sanitized = fromUnsanitized(job);
    jobUpdateController.assertNotUpdating(JOB_KEY);
    storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active());
    expectNoCronJob();
    expectInstanceQuotaCheck(sanitized, NOT_ENOUGH_QUOTA);

    control.replay();

    assertResponse(INVALID_REQUEST, thrift.createJob(job.newBuilder()));
    assertEquals(0L, statsProvider.getLongValue(CREATE_JOB));
  }

  private void assertMessageMatches(Response response, String string) {
    // TODO(wfarner): This test coverage could be much better.  Circle back to apply more thorough
    // response contents testing throughout.
    assertTrue(Iterables.any(response.getDetails(), detail -> detail.getMessage().equals(string)));
  }

  @Test
  public void testCreateEmptyJob() throws Exception {
    control.replay();

    JobConfiguration job =
        new JobConfiguration().setKey(JOB_KEY.newBuilder()).setOwner(IDENTITY);
    assertResponse(INVALID_REQUEST, thrift.createJob(job));
    assertEquals(0L, statsProvider.getLongValue(CREATE_JOB));
  }

  @Test
  public void testCreateJobFailsNoExecutorOrContainerConfig() throws Exception {
    JobConfiguration job = makeJob();
    job.getTaskConfig().unsetExecutorConfig();

    control.replay();

    Response response = thrift.createJob(job);
    assertResponse(INVALID_REQUEST, response);
    assertMessageMatches(response, ConfigurationManager.NO_EXECUTOR_OR_CONTAINER);
    assertEquals(0L, statsProvider.getLongValue(CREATE_JOB));
  }

  @Test
  public void testCreateHomogeneousJobNoInstances() throws Exception {
    JobConfiguration job = makeJob();
    job.setInstanceCount(0);

    control.replay();

    assertResponse(INVALID_REQUEST, thrift.createJob(job));
    assertEquals(0L, statsProvider.getLongValue(CREATE_JOB));
  }

  @Test
  public void testCreateJobNegativeInstanceCount() throws Exception {
    JobConfiguration job = makeJob();
    job.setInstanceCount(-1);

    control.replay();

    assertResponse(INVALID_REQUEST, thrift.createJob(job));
    assertEquals(0L, statsProvider.getLongValue(CREATE_JOB));
  }

  @Test
  public void testCreateJobNoResources() throws Exception {
    control.replay();

    TaskConfig task = productionTask();
    task.unsetResources();
    assertResponse(INVALID_REQUEST, thrift.createJob(makeJob(task)));
    assertEquals(0L, statsProvider.getLongValue(CREATE_JOB));
  }

  @Test
  public void testCreateJobBadCpu() throws Exception {
    control.replay();

    TaskConfig task = productionTask()
            .setResources(ImmutableSet.of(
                    numCpus(0.0),
                    ramMb(1024),
                    diskMb(1024)));
    assertResponse(INVALID_REQUEST, thrift.createJob(makeJob(task)));
    assertEquals(0L, statsProvider.getLongValue(CREATE_JOB));
  }

  @Test
  public void testCreateJobBadRam() throws Exception {
    control.replay();

    TaskConfig task = productionTask()
            .setResources(ImmutableSet.of(
                    numCpus(1),
                    ramMb(-123),
                    diskMb(1024)));
    assertResponse(INVALID_REQUEST, thrift.createJob(makeJob(task)));
    assertEquals(0L, statsProvider.getLongValue(CREATE_JOB));
  }

  @Test
  public void testCreateJobBadDisk() throws Exception {
    control.replay();

    TaskConfig task = productionTask()
            .setResources(ImmutableSet.of(
                    numCpus(1),
                    ramMb(1024),
                    diskMb(0)));
    task.unsetResources();
    assertResponse(INVALID_REQUEST, thrift.createJob(makeJob(task)));
  }

  @Test
  public void testCreateJobGoodResources() throws Exception {

    TaskConfig task = productionTask()
            .setResources(ImmutableSet.of(
                    numCpus(1.0),
                    ramMb(1024),
                    diskMb(1024)));

    IJobConfiguration job = IJobConfiguration.build(makeJob(task));
    SanitizedConfiguration sanitized = fromUnsanitized(job);
    jobUpdateController.assertNotUpdating(JOB_KEY);
    storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active());
    expectNoCronJob();
    expectInstanceQuotaCheck(sanitized, ENOUGH_QUOTA);

    stateManager.insertPendingTasks(
            storageUtil.mutableStoreProvider,
            sanitized.getJobConfig().getTaskConfig(),
            sanitized.getInstanceIds());

    control.replay();

    assertOkResponse(thrift.createJob(job.newBuilder()));
  }

  @Test
  public void testCreateJobPopulateDefaults() throws Exception {
    Set<Resource> resources = ImmutableSet.of(numCpus(1.0), ramMb(1024), diskMb(1024));
    TaskConfig task = new TaskConfig()
        .setContactEmail("testing@twitter.com")
        .setExecutorConfig(
            new ExecutorConfig(EXECUTOR_NAME, "config")) // Arbitrary opaque data.
        .setResources(resources)
        .setIsService(true)
        .setProduction(true)
        .setTier(TaskTestUtil.PROD_TIER_NAME)
        .setOwner(IDENTITY)
        .setContainer(Container.mesos(new MesosContainer()))
        .setJob(JOB_KEY.newBuilder());
    JobConfiguration job = makeJob(task);

    JobConfiguration sanitized = job.deepCopy();
    sanitized.getTaskConfig()
        .setJob(JOB_KEY.newBuilder())
        .setPriority(0)
        .setIsService(true)
        .setProduction(true)
        .setTaskLinks(ImmutableMap.of())
        .setConstraints(ImmutableSet.of())
        .setMaxTaskFailures(0)
        .setResources(resources);

    jobUpdateController.assertNotUpdating(JOB_KEY);
    storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active());
    expectNoCronJob();
    expectInstanceQuotaCheck(ITaskConfig.build(sanitized.getTaskConfig()), ENOUGH_QUOTA);
    stateManager.insertPendingTasks(
        storageUtil.mutableStoreProvider,
        ITaskConfig.build(sanitized.getTaskConfig()),
        ImmutableSet.of(0));

    control.replay();

    assertOkResponse(thrift.createJob(job));
    assertEquals(1L, statsProvider.getLongValue(CREATE_JOB));
  }

  @Test
  public void testCreateUnauthorizedDedicatedJob() throws Exception {
    control.replay();

    TaskConfig task = nonProductionTask();
    task.setConstraints(ImmutableSet.of(dedicatedConstraint(ImmutableSet.of("mesos"))));
    assertResponse(INVALID_REQUEST, thrift.createJob(makeJob(task)));
    assertEquals(0L, statsProvider.getLongValue(CREATE_JOB));
  }

  @Test
  public void testLimitConstraintForDedicatedJob() throws Exception {
    control.replay();

    TaskConfig task = nonProductionTask();
    task.setConstraints(ImmutableSet.of(dedicatedConstraint(1)));
    assertResponse(INVALID_REQUEST, thrift.createJob(makeJob(task)));
    assertEquals(0L, statsProvider.getLongValue(CREATE_JOB));
  }

  @Test
  public void testMultipleValueConstraintForDedicatedJob() throws Exception {
    control.replay();

    TaskConfig task = nonProductionTask();
    task.setConstraints(ImmutableSet.of(dedicatedConstraint(ImmutableSet.of("mesos", "test"))));
    assertResponse(INVALID_REQUEST, thrift.createJob(makeJob(task)));
    assertEquals(0L, statsProvider.getLongValue(CREATE_JOB));
  }

  private IScheduledTask buildTaskForJobUpdate(int instanceId) {
    return buildTaskForJobUpdate(instanceId, "data");
  }

  private IScheduledTask buildTaskForJobUpdate(int instanceId, String executorData) {
    return buildTaskForJobUpdate(
        instanceId,
        executorData,
        Optional.of(SlaPolicy.coordinatorSlaPolicy(
            new CoordinatorSlaPolicy("test-url", "test-key"))));
  }

  private IScheduledTask buildTaskForJobUpdate(int instanceId,
                                               String executorData,
                                               Optional<SlaPolicy> slaPolicy) {

    return IScheduledTask.build(new ScheduledTask()
        .setAssignedTask(new AssignedTask()
            .setInstanceId(instanceId)
            .setTask(populatedTask()
                .setIsService(true)
                .setSlaPolicy(slaPolicy.orElse(null))
                .setExecutorConfig(new ExecutorConfig().setName(EXECUTOR_NAME)
                    .setData(executorData)))));
  }

  private IScheduledTask buildScheduledTask() {
    return buildScheduledTask(JOB_NAME, TASK_ID);
  }

  private static IScheduledTask buildScheduledTask(String jobName, String taskId) {
    return IScheduledTask.build(new ScheduledTask()
        .setAssignedTask(new AssignedTask()
            .setTaskId(taskId)
            .setInstanceId(0)
            .setTask(new TaskConfig()
                .setJob(JOB_KEY.newBuilder().setName(jobName))
                .setOwner(IDENTITY))));
  }

  private void expectTransitionsToKilling() {
    expectTransitionsToKilling(Optional.empty());
  }

  private void expectTransitionsToKilling(Optional<String> message) {
    expect(auditMessages.killedByRemoteUser(message)).andReturn(Optional.of("test"));
    expect(stateManager.changeState(
        storageUtil.mutableStoreProvider,
        TASK_ID,
        Optional.empty(),
        ScheduleStatus.KILLING,
        Optional.of("test"))).andReturn(StateChangeResult.SUCCESS);
  }

  @Test
  public void testJobScopedKillsActive() throws Exception {
    Query.Builder query = Query.unscoped().byJob(JOB_KEY).active();
    storageUtil.expectTaskFetch(query, buildScheduledTask());
    jobUpdateController.assertNotUpdating(JOB_KEY);
    expectTransitionsToKilling();

    control.replay();

    assertOkResponse(thrift.killTasks(JOB_KEY.newBuilder(), null, null));
    assertEquals(1L, statsProvider.getLongValue(KILL_TASKS));
  }

  @Test
  public void testKillTasksInstanceScoped() throws Exception {
    Query.Builder query = Query.instanceScoped(JOB_KEY, ImmutableSet.of(1)).active();
    storageUtil.expectTaskFetch(query, buildScheduledTask());
    jobUpdateController.assertNotUpdating(JOB_KEY);
    expectTransitionsToKilling();

    control.replay();

    assertOkResponse(thrift.killTasks(JOB_KEY.newBuilder(), ImmutableSet.of(1), null));
    assertEquals(1L, statsProvider.getLongValue(KILL_TASKS));
  }

  @Test
  public void testKillTasksWhileJobUpdating() throws Exception {
    jobUpdateController.assertNotUpdating(JOB_KEY);
    expectLastCall().andThrow(new JobUpdatingException("job is updating"));

    control.replay();

    assertResponse(JOB_UPDATING_ERROR, thrift.killTasks(JOB_KEY.newBuilder(), null, null));
    assertEquals(0L, statsProvider.getLongValue(KILL_TASKS));
  }

  @Test
  public void testKillNonExistentTasks() throws Exception {
    Query.Builder query = Query.unscoped().byJob(JOB_KEY).active();
    storageUtil.expectTaskFetch(query);
    jobUpdateController.assertNotUpdating(JOB_KEY);

    control.replay();

    Response response = thrift.killTasks(JOB_KEY.newBuilder(), null, null);
    assertOkResponse(response);
    assertMessageMatches(response, SchedulerThriftInterface.NO_TASKS_TO_KILL_MESSAGE);
    assertEquals(0L, statsProvider.getLongValue(KILL_TASKS));
  }

  @Test
  public void testKillTasksWithMessage() throws Exception {
    String message = "Test message";
    Query.Builder query = Query.unscoped().byJob(JOB_KEY).active();
    storageUtil.expectTaskFetch(query, buildScheduledTask());
    jobUpdateController.assertNotUpdating(JOB_KEY);
    expectTransitionsToKilling(Optional.of(message));

    control.replay();

    assertOkResponse(thrift.killTasks(JOB_KEY.newBuilder(), null, message));
  }

  @Test
  public void testPruneTasksRejectsActiveStates() throws Exception {
    control.replay();
    Response rsp = thrift.pruneTasks(new TaskQuery().setStatuses(Tasks.ACTIVE_STATES));
    assertResponse(ResponseCode.ERROR, rsp);
    assertEquals(0L, statsProvider.getLongValue(PRUNE_TASKS));
  }

  @Test
  public void testPruneTasksRejectsMixedStates() throws Exception {
    control.replay();
    Response rsp = thrift.pruneTasks(new TaskQuery().setStatuses(
        ImmutableSet.of(ScheduleStatus.FINISHED, ScheduleStatus.KILLING)));
    assertResponse(ResponseCode.ERROR, rsp);
    assertEquals(0L, statsProvider.getLongValue(PRUNE_TASKS));
  }

  @Test
  public void testPruneTasksAddsDefaultStatuses() throws Exception {
    storageUtil.expectTaskFetch(
        Query.arbitrary(new TaskQuery().setStatuses(Tasks.TERMINAL_STATES)),
        buildScheduledTask());
    stateManager.deleteTasks(
        storageUtil.mutableStoreProvider,
        ImmutableSet.of(buildScheduledTask().getAssignedTask().getTaskId()));
    control.replay();

    assertOkResponse(thrift.pruneTasks(new TaskQuery()));
    assertEquals(1L, statsProvider.getLongValue(PRUNE_TASKS));
  }

  @Test
  public void testPruneTasksRespectsScopedTerminalState() throws Exception {
    storageUtil.expectTaskFetch(
        Query.arbitrary(new TaskQuery().setStatuses(ImmutableSet.of(ScheduleStatus.FAILED))),
        buildScheduledTask());
    stateManager.deleteTasks(
        storageUtil.mutableStoreProvider,
        ImmutableSet.of(buildScheduledTask().getAssignedTask().getTaskId()));
    control.replay();

    assertOkResponse(thrift.pruneTasks(
        new TaskQuery().setStatuses(ImmutableSet.of(ScheduleStatus.FAILED))));
    assertEquals(1L, statsProvider.getLongValue(PRUNE_TASKS));
  }

  @Test
  public void testPruneTasksAppliesQueryLimit() throws Exception {
    TaskQuery query = new TaskQuery().setLimit(3);
    storageUtil.expectTaskFetch(
        Query.arbitrary(query.setStatuses(Tasks.TERMINAL_STATES)),
        buildScheduledTask("a/b/c", "task1"),
        buildScheduledTask("a/b/c", "task2"),
        buildScheduledTask("a/b/c", "task3"),
        buildScheduledTask("a/b/c", "task4"),
        buildScheduledTask("a/b/c", "task5"));
    stateManager.deleteTasks(
        storageUtil.mutableStoreProvider,
        ImmutableSet.of("task1", "task2", "task3"));
    control.replay();

    assertOkResponse(thrift.pruneTasks(query));
    assertEquals(3L, statsProvider.getLongValue(PRUNE_TASKS));
  }

  @Test
  public void testSetQuota() throws Exception {
    ResourceAggregate resourceAggregate = new ResourceAggregate()
        .setResources(ImmutableSet.of(numCpus(10.0), ramMb(200), diskMb(100)));
    quotaManager.saveQuota(
        ROLE,
        IResourceAggregate.build(resourceAggregate),
        storageUtil.mutableStoreProvider);

    control.replay();

    assertOkResponse(thrift.setQuota(ROLE, resourceAggregate));
  }

  @Test
  public void testSetQuotaFails() throws Exception {
    ResourceAggregate resourceAggregate = new ResourceAggregate()
        .setResources(ImmutableSet.of(numCpus(10.0), ramMb(200), diskMb(100)));
    quotaManager.saveQuota(
        ROLE,
        IResourceAggregate.build(resourceAggregate.deepCopy()
            .setResources(ImmutableSet.of(numCpus(10), ramMb(200), diskMb(100)))),
        storageUtil.mutableStoreProvider);

    expectLastCall().andThrow(new QuotaManager.QuotaException("fail"));

    control.replay();

    assertResponse(INVALID_REQUEST, thrift.setQuota(ROLE, resourceAggregate));
  }

  @Test
  public void testForceTaskState() throws Exception {
    ScheduleStatus status = ScheduleStatus.FAILED;

    expect(auditMessages.transitionedBy()).andReturn(Optional.of("test"));
    expect(stateManager.changeState(
        storageUtil.mutableStoreProvider,
        TASK_ID,
        Optional.empty(),
        ScheduleStatus.FAILED,
        Optional.of("test"))).andReturn(StateChangeResult.SUCCESS);

    control.replay();

    assertOkResponse(thrift.forceTaskState(TASK_ID, status));
  }

  @Test
  public void testBackupControls() throws Exception {
    backup.backupNow();

    Set<String> backups = ImmutableSet.of("a", "b");
    expect(recovery.listBackups()).andReturn(backups);

    String backupId = "backup";
    recovery.stage(backupId);

    Query.Builder query = Query.taskScoped("taskId");
    Set<IScheduledTask> queryResult = ImmutableSet.of(
        IScheduledTask.build(new ScheduledTask().setStatus(ScheduleStatus.RUNNING)));
    expect(recovery.query(query)).andReturn(queryResult);

    recovery.deleteTasks(query);

    recovery.commit();

    recovery.unload();

    control.replay();

    assertEquals(okEmptyResponse(), thrift.performBackup());

    assertEquals(
        okResponse(Result.listBackupsResult(new ListBackupsResult().setBackups(backups))),
        thrift.listBackups());

    assertEquals(okEmptyResponse(), thrift.stageRecovery(backupId));

    assertEquals(
        okResponse(Result.queryRecoveryResult(
            new QueryRecoveryResult().setTasks(IScheduledTask.toBuildersSet(queryResult)))),
        thrift.queryRecovery(query.get().newBuilder()));

    assertEquals(okEmptyResponse(), thrift.deleteRecoveryTasks(query.get().newBuilder()));

    assertEquals(okEmptyResponse(), thrift.commitRecovery());

    assertEquals(okEmptyResponse(), thrift.unloadRecovery());
  }

  @Test
  public void testRecoveryException() throws Exception {
    Throwable recoveryException = new RecoveryException("Injected");

    String backupId = "backup";
    recovery.stage(backupId);
    expectLastCall().andThrow(recoveryException);

    control.replay();

    try {
      thrift.stageRecovery(backupId);
      fail("No recovery exception thrown.");
    } catch (RecoveryException e) {
      assertEquals(recoveryException.getMessage(), e.getMessage());
    }
  }

  @Test
  public void testRestartShards() throws Exception {
    Set<Integer> shards = ImmutableSet.of(0);

    jobUpdateController.assertNotUpdating(JOB_KEY);
    storageUtil.expectTaskFetch(
        Query.instanceScoped(JOB_KEY, shards).active(),
        buildScheduledTask());

    expect(auditMessages.restartedByRemoteUser())
        .andReturn(Optional.of("test"));
    expect(stateManager.changeState(
        storageUtil.mutableStoreProvider,
        TASK_ID,
        Optional.empty(),
        ScheduleStatus.RESTARTING,
        Optional.of("test"))).andReturn(StateChangeResult.SUCCESS);

    control.replay();

    assertOkResponse(
        thrift.restartShards(JOB_KEY.newBuilder(), shards));
    assertEquals(1L, statsProvider.getLongValue(RESTART_SHARDS));
  }

  @Test
  public void testRestartShardsWhileJobUpdating() throws Exception {
    Set<Integer> shards = ImmutableSet.of(1, 6);

    jobUpdateController.assertNotUpdating(JOB_KEY);
    expectLastCall().andThrow(new JobUpdatingException("job is updating"));

    control.replay();

    assertResponse(
        JOB_UPDATING_ERROR,
        thrift.restartShards(JOB_KEY.newBuilder(), shards));
    assertEquals(0L, statsProvider.getLongValue(KILL_TASKS));
  }

  @Test
  public void testRestartShardsNotFoundTasksFailure() throws Exception {
    Set<Integer> shards = ImmutableSet.of(1, 6);

    jobUpdateController.assertNotUpdating(JOB_KEY);
    storageUtil.expectTaskFetch(Query.instanceScoped(JOB_KEY, shards).active());

    control.replay();

    assertResponse(
        INVALID_REQUEST,
        thrift.restartShards(JOB_KEY.newBuilder(), shards));
    assertEquals(0L, statsProvider.getLongValue(KILL_TASKS));
  }

  @Test
  public void testReplaceCronTemplate() throws Exception {
    jobUpdateController.assertNotUpdating(JOB_KEY);
    SanitizedConfiguration sanitized = fromUnsanitized(IJobConfiguration.build(CRON_JOB));

    expectCronQuotaCheck(sanitized.getJobConfig(), ENOUGH_QUOTA);
    cronJobManager.updateJob(anyObject(SanitizedCronJob.class));
    control.replay();

    assertOkResponse(thrift.replaceCronTemplate(CRON_JOB));
    assertEquals(1L, statsProvider.getLongValue(CREATE_OR_UPDATE_CRON));
  }

  @Test
  public void testReplaceCronTemplateDoesNotExist() throws Exception {
    jobUpdateController.assertNotUpdating(JOB_KEY);
    SanitizedConfiguration sanitized = fromUnsanitized(IJobConfiguration.build(CRON_JOB));

    expectCronQuotaCheck(sanitized.getJobConfig(), ENOUGH_QUOTA);
    cronJobManager.updateJob(anyObject(SanitizedCronJob.class));
    expectLastCall().andThrow(new CronException("Nope"));

    control.replay();

    assertResponse(INVALID_REQUEST, thrift.replaceCronTemplate(CRON_JOB));
    assertEquals(0L, statsProvider.getLongValue(CREATE_OR_UPDATE_CRON));
  }

  @Test
  public void testStartCronJob() throws Exception {
    cronJobManager.startJobNow(JOB_KEY);
    control.replay();
    assertResponse(OK, thrift.startCronJob(JOB_KEY.newBuilder()));
  }

  @Test
  public void testStartCronJobFailsInCronManager() throws Exception {
    cronJobManager.startJobNow(JOB_KEY);
    expectLastCall().andThrow(new CronException("failed"));
    control.replay();
    assertResponse(INVALID_REQUEST, thrift.startCronJob(JOB_KEY.newBuilder()));
  }

  @Test
  public void testScheduleCronCreatesJob() throws Exception {
    jobUpdateController.assertNotUpdating(JOB_KEY);
    SanitizedConfiguration sanitized = fromUnsanitized(IJobConfiguration.build(CRON_JOB));

    expectCronQuotaCheck(sanitized.getJobConfig(), ENOUGH_QUOTA);

    expectNoCronJob().times(2);
    storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active());
    cronJobManager.createJob(SanitizedCronJob.from(sanitized));
    control.replay();
    assertResponse(OK, thrift.scheduleCronJob(CRON_JOB));
  }

  @Test
  public void testScheduleCronFailsCreationDueToExistingNonCron() throws Exception {
    jobUpdateController.assertNotUpdating(JOB_KEY);
    SanitizedConfiguration sanitized = fromUnsanitized(IJobConfiguration.build(CRON_JOB));

    expectCronQuotaCheck(sanitized.getJobConfig(), ENOUGH_QUOTA);

    expectNoCronJob();
    storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active(), buildScheduledTask());
    control.replay();
    assertEquals(
        invalidResponse(jobAlreadyExistsMessage(JOB_KEY)),
        thrift.scheduleCronJob(CRON_JOB));
  }

  @Test
  public void testScheduleCronUpdatesJob() throws Exception {
    jobUpdateController.assertNotUpdating(JOB_KEY);
    SanitizedConfiguration sanitized = fromUnsanitized(IJobConfiguration.build(CRON_JOB));

    expectCronQuotaCheck(sanitized.getJobConfig(), ENOUGH_QUOTA);

    expectCronJob();
    cronJobManager.updateJob(SanitizedCronJob.from(sanitized));
    control.replay();

    assertResponse(OK, thrift.scheduleCronJob(CRON_JOB));
  }

  @Test
  public void testScheduleCronJobFailedTaskConfigValidation() throws Exception {
    control.replay();
    IJobConfiguration job = IJobConfiguration.build(makeJob(INVALID_TASK_CONFIG));
    assertResponse(INVALID_REQUEST, thrift.scheduleCronJob(job.newBuilder()));
  }

  @Test
  public void testScheduleCronJobFailsWithNoCronSchedule() throws Exception {
    control.replay();

    assertEquals(
        invalidResponse(noCronScheduleMessage(JOB_KEY)),
        thrift.scheduleCronJob(makeJob()));
  }

  @Test
  public void testScheduleCronFailsQuotaCheck() throws Exception {
    jobUpdateController.assertNotUpdating(JOB_KEY);
    SanitizedConfiguration sanitized = fromUnsanitized(IJobConfiguration.build(CRON_JOB));

    expectCronQuotaCheck(sanitized.getJobConfig(), NOT_ENOUGH_QUOTA);

    control.replay();
    assertResponse(INVALID_REQUEST, thrift.scheduleCronJob(CRON_JOB));
  }

  @Test
  public void testDescheduleCronJob() throws Exception {
    jobUpdateController.assertNotUpdating(JOB_KEY);
    expect(cronJobManager.deleteJob(JOB_KEY)).andReturn(true);

    control.replay();

    assertResponse(OK, thrift.descheduleCronJob(CRON_JOB.getKey()));
  }

  @Test
  public void testDescheduleCronJobWhileJobUpdating() throws Exception {
    jobUpdateController.assertNotUpdating(JOB_KEY);
    expectLastCall().andThrow(new JobUpdatingException("job is updating"));
    control.replay();
    assertResponse(JOB_UPDATING_ERROR, thrift.descheduleCronJob(CRON_JOB.getKey()));
  }

  @Test
  public void testDescheduleNotACron() throws Exception {
    jobUpdateController.assertNotUpdating(JOB_KEY);
    expect(cronJobManager.deleteJob(JOB_KEY)).andReturn(false);
    control.replay();

    assertEquals(
        okEmptyResponse(notScheduledCronMessage(JOB_KEY)),
        thrift.descheduleCronJob(JOB_KEY.newBuilder()));
  }

  @Test
  public void testUnauthorizedDedicatedJob() throws Exception {
    control.replay();

    TaskConfig task = nonProductionTask();
    task.setConstraints(ImmutableSet.of(dedicatedConstraint(ImmutableSet.of("mesos"))));
    assertResponse(INVALID_REQUEST, thrift.createJob(makeJob(task)));
    assertEquals(0L, statsProvider.getLongValue(CREATE_JOB));
  }

  private static Set<IHostStatus> status(String host, MaintenanceMode mode) {
    return ImmutableSet.of(IHostStatus.build(new HostStatus(host, mode)));
  }

  @Test
  public void testHostMaintenance() throws Exception {
    Set<String> hostnames = ImmutableSet.of("a");
    Set<IHostStatus> none = status("a", NONE);
    Set<IHostStatus> scheduled = status("a", SCHEDULED);
    Set<IHostStatus> draining = status("a", DRAINING);
    Set<IHostStatus> drained = status("a", DRAINING);
    expect(maintenance.getStatus(hostnames)).andReturn(none);
    expect(maintenance.startMaintenance(hostnames)).andReturn(scheduled);
    expect(maintenance.drain(hostnames)).andReturn(draining);
    expect(maintenance.getStatus(hostnames)).andReturn(draining);
    expect(maintenance.getStatus(hostnames)).andReturn(drained);
    expect(maintenance.endMaintenance(hostnames)).andReturn(none);

    control.replay();

    Hosts hosts = new Hosts(hostnames);

    assertEquals(
        IHostStatus.toBuildersSet(none),
        thrift.maintenanceStatus(hosts).getResult().getMaintenanceStatusResult()
            .getStatuses());
    assertEquals(1L, statsProvider.getLongValue(MAINTENANCE_STATUS));
    assertEquals(
        IHostStatus.toBuildersSet(scheduled),
        thrift.startMaintenance(hosts).getResult().getStartMaintenanceResult()
            .getStatuses());
    assertEquals(1L, statsProvider.getLongValue(START_MAINTENANCE));
    assertEquals(
        IHostStatus.toBuildersSet(draining),
        thrift.drainHosts(hosts).getResult().getDrainHostsResult().getStatuses());
    assertEquals(1L, statsProvider.getLongValue(DRAIN_HOSTS));
    assertEquals(
        IHostStatus.toBuildersSet(draining),
        thrift.maintenanceStatus(hosts).getResult().getMaintenanceStatusResult()
            .getStatuses());
    assertEquals(2L, statsProvider.getLongValue(MAINTENANCE_STATUS));
    assertEquals(
        IHostStatus.toBuildersSet(drained),
        thrift.maintenanceStatus(hosts).getResult().getMaintenanceStatusResult()
            .getStatuses());
    assertEquals(3L, statsProvider.getLongValue(MAINTENANCE_STATUS));
    assertEquals(
        IHostStatus.toBuildersSet(none),
        thrift.endMaintenance(hosts).getResult().getEndMaintenanceResult().getStatuses());
    assertEquals(1L, statsProvider.getLongValue(END_MAINTENANCE));
  }

  @Test
  public void testSLAHostMaintenance() throws Exception {
    Set<String> hostnames = ImmutableSet.of("a");
    SlaPolicy defaultSlaPolicy = SlaPolicy.percentageSlaPolicy(
        new PercentageSlaPolicy().setPercentage(95));
    Set<IHostStatus> none = status("a", NONE);
    Set<IHostStatus> scheduled = status("a", SCHEDULED);
    Set<IHostStatus> draining = status("a", DRAINING);
    Set<IHostStatus> drained = status("a", DRAINING);
    expect(maintenance.getStatus(hostnames)).andReturn(none);
    expect(maintenance.startMaintenance(hostnames)).andReturn(scheduled);
    expect(maintenance.slaDrain(hostnames, defaultSlaPolicy, 10)).andReturn(draining);
    expect(maintenance.getStatus(hostnames)).andReturn(draining);
    expect(maintenance.getStatus(hostnames)).andReturn(drained);
    expect(maintenance.endMaintenance(hostnames)).andReturn(none);

    control.replay();

    Hosts hosts = new Hosts(hostnames);

    assertEquals(
        IHostStatus.toBuildersSet(none),
        thrift.maintenanceStatus(hosts).getResult().getMaintenanceStatusResult()
            .getStatuses());
    assertEquals(1L, statsProvider.getLongValue(MAINTENANCE_STATUS));
    assertEquals(
        IHostStatus.toBuildersSet(scheduled),
        thrift.startMaintenance(hosts).getResult().getStartMaintenanceResult()
            .getStatuses());
    assertEquals(1L, statsProvider.getLongValue(START_MAINTENANCE));
    assertEquals(
        IHostStatus.toBuildersSet(draining),
        thrift.slaDrainHosts(hosts, defaultSlaPolicy, 10)
            .getResult()
            .getDrainHostsResult()
            .getStatuses());
    assertEquals(1L, statsProvider.getLongValue(SLA_DRAIN_HOSTS));
    assertEquals(
        IHostStatus.toBuildersSet(draining),
        thrift.maintenanceStatus(hosts).getResult().getMaintenanceStatusResult()
            .getStatuses());
    assertEquals(2L, statsProvider.getLongValue(MAINTENANCE_STATUS));
    assertEquals(
        IHostStatus.toBuildersSet(drained),
        thrift.maintenanceStatus(hosts).getResult().getMaintenanceStatusResult()
            .getStatuses());
    assertEquals(3L, statsProvider.getLongValue(MAINTENANCE_STATUS));
    assertEquals(
        IHostStatus.toBuildersSet(none),
        thrift.endMaintenance(hosts).getResult().getEndMaintenanceResult().getStatuses());
    assertEquals(1L, statsProvider.getLongValue(END_MAINTENANCE));
  }

  private static Response okEmptyResponse() {
    return response(OK, Optional.empty());
  }

  private static Response okEmptyResponse(String message) {
    return Responses.empty()
        .setResponseCode(OK)
        .setDetails(ImmutableList.of(new ResponseDetail(message)));
  }

  private static Response invalidResponse(String message) {
    return Responses.empty()
        .setResponseCode(INVALID_REQUEST)
        .setDetails(ImmutableList.of(new ResponseDetail(message)));
  }

  @Test
  public void testSnapshot() throws Exception {
    snapshotStore.snapshot();
    expectLastCall();

    snapshotStore.snapshot();
    expectLastCall().andThrow(new StorageException("mock error!"));

    control.replay();

    assertOkResponse(thrift.snapshot());

    try {
      thrift.snapshot();
      fail("No StorageException thrown.");
    } catch (StorageException e) {
      // Expected.
    }
  }

  @Test
  public void testExplicitTaskReconciliationWithNoBatchSize() throws Exception {
    ExplicitReconciliationSettings settings = new ExplicitReconciliationSettings();

    taskReconciler.triggerExplicitReconciliation(Optional.empty());
    expectLastCall();

    control.replay();

    assertOkResponse(thrift.triggerExplicitTaskReconciliation(settings));
  }

  @Test
  public void testExplicitTaskReconciliationWithValidBatchSize() throws Exception {
    ExplicitReconciliationSettings settings = new ExplicitReconciliationSettings();
    settings.setBatchSize(10);

    taskReconciler.triggerExplicitReconciliation(Optional.of(settings.getBatchSize()));
    expectLastCall();

    control.replay();
    assertOkResponse(thrift.triggerExplicitTaskReconciliation(settings));
  }

  @Test
  public void testExplicitTaskReconciliationWithNegativeBatchSize() throws Exception {
    ExplicitReconciliationSettings settings = new ExplicitReconciliationSettings();
    settings.setBatchSize(-1000);

    control.replay();
    assertResponse(INVALID_REQUEST, thrift.triggerExplicitTaskReconciliation(settings));
  }

  @Test
  public void testExplicitTaskReconciliationWithZeroBatchSize() throws Exception {
    ExplicitReconciliationSettings settings = new ExplicitReconciliationSettings();
    settings.setBatchSize(0);

    control.replay();
    assertResponse(INVALID_REQUEST, thrift.triggerExplicitTaskReconciliation(settings));
  }

  @Test
  public void testImplicitTaskReconciliation() throws Exception {
    taskReconciler.triggerImplicitReconciliation();
    expectLastCall();

    control.replay();

    assertOkResponse(thrift.triggerImplicitTaskReconciliation());
  }

  @Test
  public void testAddInstancesWithInstanceKey() throws Exception {
    expectNoCronJob();
    jobUpdateController.assertNotUpdating(JOB_KEY);
    IScheduledTask activeTask = buildScheduledTask();
    ITaskConfig task = activeTask.getAssignedTask().getTask();
    storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active(), activeTask);
    expect(quotaManager.checkInstanceAddition(
        task,
        2,
        storageUtil.mutableStoreProvider)).andReturn(ENOUGH_QUOTA);
    stateManager.insertPendingTasks(
        storageUtil.mutableStoreProvider,
        task,
        ImmutableSet.of(1, 2));

    control.replay();

    assertOkResponse(thrift.addInstances(INSTANCE_KEY, 2));
    assertEquals(2L, statsProvider.getLongValue(ADD_INSTANCES));
  }

  @Test
  public void testAddInstancesWithInstanceKeyFailsWithNoInstance() throws Exception {
    expectNoCronJob();
    jobUpdateController.assertNotUpdating(JOB_KEY);
    storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active());

    control.replay();

    assertEquals(
        invalidResponse(SchedulerThriftInterface.INVALID_INSTANCE_ID),
        thrift.addInstances(INSTANCE_KEY, 2));
    assertEquals(0L, statsProvider.getLongValue(ADD_INSTANCES));
  }

  @Test
  public void testAddInstancesWithInstanceKeyFailsInvalidCount() throws Exception {
    control.replay();

    assertEquals(
        invalidResponse(SchedulerThriftInterface.INVALID_INSTANCE_COUNT),
        thrift.addInstances(INSTANCE_KEY, 0));
    assertEquals(0L, statsProvider.getLongValue(ADD_INSTANCES));
  }

  @Test
  public void testAddInstancesFailsCronJob() throws Exception {
    expectCronJob();

    control.replay();

    assertResponse(INVALID_REQUEST, thrift.addInstances(INSTANCE_KEY, 1));
  }

  @Test(expected = StorageException.class)
  public void testAddInstancesFailsWithNonTransient() throws Exception {
    expect(storageUtil.jobStore.fetchJob(JOB_KEY)).andThrow(new StorageException("no retry"));

    control.replay();

    thrift.addInstances(INSTANCE_KEY, 1);
  }

  @Test
  public void testAddInstancesWhileJobUpdating() throws Exception {
    expectNoCronJob();
    jobUpdateController.assertNotUpdating(JOB_KEY);
    expectLastCall().andThrow(new JobUpdatingException("job is updating"));

    control.replay();

    assertResponse(JOB_UPDATING_ERROR, thrift.addInstances(INSTANCE_KEY, 1));
    assertEquals(0L, statsProvider.getLongValue(ADD_INSTANCES));
  }

  @Test
  public void testAddInstancesFailsQuotaCheck() throws Exception {
    IScheduledTask activeTask = buildScheduledTask();
    ITaskConfig task = activeTask.getAssignedTask().getTask();
    expectNoCronJob();
    jobUpdateController.assertNotUpdating(JOB_KEY);
    storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active(), activeTask);
    expectInstanceQuotaCheck(task, NOT_ENOUGH_QUOTA);

    control.replay();

    assertResponse(INVALID_REQUEST, thrift.addInstances(INSTANCE_KEY, 1));
    assertEquals(0L, statsProvider.getLongValue(ADD_INSTANCES));
  }

  @Test
  public void testAddInstancesInstanceCollisionFailure() throws Exception {
    IScheduledTask activeTask = buildScheduledTask();
    ITaskConfig task = activeTask.getAssignedTask().getTask();
    expectNoCronJob();
    jobUpdateController.assertNotUpdating(JOB_KEY);
    storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active(), activeTask);
    expectInstanceQuotaCheck(task, ENOUGH_QUOTA);
    stateManager.insertPendingTasks(
        storageUtil.mutableStoreProvider,
        task,
        ImmutableSet.of(1));
    expectLastCall().andThrow(new IllegalArgumentException("instance collision"));

    control.replay();

    assertResponse(INVALID_REQUEST, thrift.addInstances(INSTANCE_KEY, 1));
    assertEquals(0L, statsProvider.getLongValue(ADD_INSTANCES));
  }

  @Test
  public void testStartUpdate() throws Exception {
    expectGetRemoteUser();
    expectNoCronJob();

    ITaskConfig newTask = buildTaskForJobUpdate(0).getAssignedTask().getTask();

    IScheduledTask oldTask1 = buildTaskForJobUpdate(0, "old");
    IScheduledTask oldTask2 = buildTaskForJobUpdate(1, "old");
    IScheduledTask oldTask3 = buildTaskForJobUpdate(2, "old2");
    IScheduledTask oldTask4 = buildTaskForJobUpdate(3, "old2");
    IScheduledTask oldTask5 = buildTaskForJobUpdate(4, "old");
    IScheduledTask oldTask6 = buildTaskForJobUpdate(5, "old");
    IScheduledTask oldTask7 = buildTaskForJobUpdate(6, "old");

    IJobUpdate update = buildJobUpdate(6, newTask, ImmutableMap.of(
        oldTask1.getAssignedTask().getTask(), ImmutableSet.of(new Range(0, 1), new Range(4, 6)),
        oldTask3.getAssignedTask().getTask(), ImmutableSet.of(new Range(2, 3))
    ));

    expect(quotaManager.checkJobUpdate(
        update,
        storageUtil.mutableStoreProvider)).andReturn(ENOUGH_QUOTA);

    expect(uuidGenerator.createNew()).andReturn(UU_ID);
    storageUtil.expectTaskFetch(
        Query.unscoped().byJob(JOB_KEY).active(),
        oldTask1,
        oldTask2,
        oldTask3,
        oldTask4,
        oldTask5,
        oldTask6,
        oldTask7);

    jobUpdateController.start(update, AUDIT);

    control.replay();

    Response response =
        assertOkResponse(thrift.startJobUpdate(buildJobUpdateRequest(update), AUDIT_MESSAGE));
    assertEquals(
        new StartJobUpdateResult(UPDATE_KEY.newBuilder()).setUpdateSummary(
        update.getSummary().newBuilder()), response.getResult().getStartJobUpdateResult());
    assertEquals(6L, statsProvider.getLongValue(START_JOB_UPDATE));
  }

  private void expectJobUpdateQuotaCheck(QuotaCheckResult result) {
    expect(quotaManager.checkJobUpdate(
        anyObject(IJobUpdate.class),
        eq(storageUtil.mutableStoreProvider))).andReturn(result);
  }

  @Test
  public void testStartUpdateEmptyDesired() throws Exception {
    expectGetRemoteUser();
    expectNoCronJob();

    ITaskConfig newTask = buildTaskForJobUpdate(0).getAssignedTask().getTask();

    IScheduledTask oldTask1 = buildTaskForJobUpdate(0);
    IScheduledTask oldTask2 = buildTaskForJobUpdate(1);

    // Set instance count to 1 to generate empty desired state in diff.
    IJobUpdate update = buildJobUpdate(
        1,
        newTask,
        ImmutableMap.of(oldTask1.getAssignedTask().getTask(), ImmutableSet.of(new Range(0, 1))));

    expectJobUpdateQuotaCheck(ENOUGH_QUOTA);

    // Set diff-adjusted IJobUpdate expectations.
    JobUpdate expected = update.newBuilder();
    expected.getInstructions().setInitialState(ImmutableSet.of(
        new InstanceTaskConfig(newTask.newBuilder(), ImmutableSet.of(new Range(1, 1)))));
    expected.getInstructions().unsetDesiredState();

    expect(uuidGenerator.createNew()).andReturn(UU_ID);
    storageUtil.expectTaskFetch(
        Query.unscoped().byJob(JOB_KEY).active(),
        oldTask1,
        oldTask2);

    jobUpdateController.start(IJobUpdate.build(expected), AUDIT);

    control.replay();

    Response response = assertOkResponse(
        thrift.startJobUpdate(buildJobUpdateRequest(update), AUDIT_MESSAGE));
    assertEquals(
        new StartJobUpdateResult(UPDATE_KEY.newBuilder()).setUpdateSummary(expected.getSummary()),
        response.getResult().getStartJobUpdateResult());
    assertEquals(1L, statsProvider.getLongValue(START_JOB_UPDATE));
  }

  @Test(expected = NullPointerException.class)
  public void testStartUpdateFailsNullRequest() throws Exception {
    control.replay();
    thrift.startJobUpdate(null, AUDIT_MESSAGE);
  }

  @Test(expected = NullPointerException.class)
  public void testStartUpdateFailsNullTaskConfig() throws Exception {
    control.replay();
    thrift.startJobUpdate(
        new JobUpdateRequest(null, 5, buildJobUpdateSettings()),
        AUDIT_MESSAGE);
  }

  @Test
  public void testStartUpdateFailsInvalidGroupSize() throws Exception {
    control.replay();

    JobUpdateRequest updateRequest = buildServiceJobUpdateRequest();
    updateRequest.getSettings().setUpdateStrategy(
        JobUpdateStrategy.queueStrategy(new QueueJobUpdateStrategy().setGroupSize(0)));

    assertEquals(
        invalidResponse(SchedulerThriftInterface.NO_INSTANCES_MODIFIED),
        thrift.startJobUpdate(updateRequest, AUDIT_MESSAGE));
    assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE));
  }

  @Test
  public void testStartUpdateFailsInvalidGroupSizeVariableBatch() throws Exception {
    control.replay();

    JobUpdateRequest updateRequest = buildServiceJobUpdateRequest();
    updateRequest.getSettings().setUpdateStrategy(
            JobUpdateStrategy.varBatchStrategy(
                    new VariableBatchJobUpdateStrategy().setGroupSizes(ImmutableList.of(1, 0, 4))));

    assertEquals(
            invalidResponse(SchedulerThriftInterface.INVALID_GROUP_SIZE),
            thrift.startJobUpdate(updateRequest, AUDIT_MESSAGE));
    assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE));
  }

  @Test
  public void testStartUpdateFailsInvalidGroupsSum() throws Exception {
    control.replay();

    JobUpdateRequest updateRequest = buildServiceJobUpdateRequest();
    updateRequest.getSettings().setUpdateStrategy(
            JobUpdateStrategy.varBatchStrategy(
                    new VariableBatchJobUpdateStrategy().setGroupSizes(ImmutableList.of())));

    assertEquals(
            invalidResponse(SchedulerThriftInterface.NO_INSTANCES_MODIFIED),
            thrift.startJobUpdate(updateRequest, AUDIT_MESSAGE));
    assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE));
  }

  @Test
  public void testStartUpdateFailsInvalidMaxInstanceFailures() throws Exception {
    control.replay();

    JobUpdateRequest updateRequest = buildServiceJobUpdateRequest();
    updateRequest.getSettings().setMaxPerInstanceFailures(-1);

    assertEquals(
        invalidResponse(SchedulerThriftInterface.INVALID_MAX_INSTANCE_FAILURES),
        thrift.startJobUpdate(updateRequest, AUDIT_MESSAGE));
    assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE));
  }

  @Test
  public void testStartUpdateFailsTooManyPerInstanceFailures() throws Exception {
    control.replay();

    JobUpdateRequest updateRequest = buildServiceJobUpdateRequest();
    updateRequest.getSettings()
        .setMaxPerInstanceFailures(THRESHOLDS.getMaxUpdateInstanceFailures() + 10);

    assertEquals(
        invalidResponse(SchedulerThriftInterface.TOO_MANY_POTENTIAL_FAILED_INSTANCES),
        thrift.startJobUpdate(updateRequest, AUDIT_MESSAGE));
    assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE));
  }

  @Test
  public void testStartUpdateFailsInvalidMaxFailedInstances() throws Exception {
    control.replay();

    JobUpdateRequest updateRequest = buildServiceJobUpdateRequest();
    updateRequest.getSettings().setMaxFailedInstances(-1);

    assertEquals(
        invalidResponse(SchedulerThriftInterface.INVALID_MAX_FAILED_INSTANCES),
        thrift.startJobUpdate(updateRequest, AUDIT_MESSAGE));
    assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE));
  }

  @Test
  public void testStartUpdateFailsInvalidMinWaitInRunning() throws Exception {
    control.replay();

    JobUpdateRequest updateRequest = buildServiceJobUpdateRequest();
    updateRequest.getSettings().setMinWaitInInstanceRunningMs(-1);

    assertEquals(
        invalidResponse(SchedulerThriftInterface.INVALID_MIN_WAIT_TO_RUNNING),
        thrift.startJobUpdate(updateRequest, AUDIT_MESSAGE));
    assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE));
  }

  @Test
  public void testStartUpdateFailsNonServiceTask() throws Exception {
    control.replay();
    JobUpdateRequest request = buildJobUpdateRequest(populatedTask().setIsService(false));
    assertResponse(INVALID_REQUEST, thrift.startJobUpdate(request, AUDIT_MESSAGE));
    assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE));
  }

  @Test
  public void testStartUpdateFailsInvalidPulseTimeout() throws Exception {
    control.replay();

    JobUpdateRequest updateRequest = buildServiceJobUpdateRequest();
    updateRequest.getSettings().setBlockIfNoPulsesAfterMs(-1);

    assertEquals(
        invalidResponse(SchedulerThriftInterface.INVALID_PULSE_TIMEOUT),
        thrift.startJobUpdate(updateRequest, AUDIT_MESSAGE));
    assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE));
  }

  @Test
  public void testStartUpdateFailsForCronJob() throws Exception {
    JobUpdateRequest request = buildServiceJobUpdateRequest(populatedTask());
    expectCronJob();

    control.replay();
    assertEquals(invalidResponse(NO_CRON), thrift.startJobUpdate(request, AUDIT_MESSAGE));
    assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE));
  }

  @Test
  public void testStartUpdateFailsConfigValidation() throws Exception {
    JobUpdateRequest request = buildJobUpdateRequest(INVALID_TASK_CONFIG.setIsService(true));

    control.replay();
    assertResponse(INVALID_REQUEST, thrift.startJobUpdate(request, AUDIT_MESSAGE));
    assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE));
  }

  @Test
  public void testStartNoopUpdate() throws Exception {
    expectNoCronJob();
    expect(uuidGenerator.createNew()).andReturn(UU_ID);
    ITaskConfig newTask = buildTaskForJobUpdate(0).getAssignedTask().getTask();

    IScheduledTask oldTask = buildTaskForJobUpdate(0);
    storageUtil.expectTaskFetch(Query.unscoped().byJob(JOB_KEY).active(), oldTask);

    IJobUpdate update = buildJobUpdate(
        1,
        newTask,
        ImmutableMap.of(oldTask.getAssignedTask().getTask(), ImmutableSet.of(new Range(0, 0))));

    control.replay();
    JobUpdateRequest request = buildJobUpdateRequest(update);
    Response response = thrift.startJobUpdate(request, AUDIT_MESSAGE);
    assertResponse(OK, response);
    assertEquals(
        NOOP_JOB_UPDATE_MESSAGE,
        Iterables.getOnlyElement(response.getDetails()).getMessage());
    assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE));
  }

  @Test
  public void testStartUpdateInvalidScope() throws Exception {
    expectNoCronJob();
    expect(uuidGenerator.createNew()).andReturn(UU_ID);

    IScheduledTask oldTask = buildTaskForJobUpdate(0);
    storageUtil.expectTaskFetch(Query.unscoped().byJob(JOB_KEY).active(), oldTask);

    ITaskConfig newTask = buildTaskForJobUpdate(0).getAssignedTask().getTask();
    JobUpdate builder = buildJobUpdate(
        1,
        newTask,
        ImmutableMap.of(oldTask.getAssignedTask().getTask(), ImmutableSet.of(new Range(0, 0))))
        .newBuilder();
    builder.getInstructions().getSettings()
        .setUpdateOnlyTheseInstances(ImmutableSet.of(new Range(100, 100)));

    control.replay();
    JobUpdateRequest request = buildJobUpdateRequest(IJobUpdate.build(builder));
    assertResponse(INVALID_REQUEST, thrift.startJobUpdate(request, AUDIT_MESSAGE));
    assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE));
  }

  @Test
  public void testStartScopeIncludesNoop() throws Exception {
    // Test for regression of AURORA-1332: a scoped update should be allowed when unchanged
    // instances are included in the scope.

    expectGetRemoteUser();
    expectNoCronJob();
    expect(uuidGenerator.createNew()).andReturn(UU_ID);

    ScheduledTask taskBuilder = buildTaskForJobUpdate(0).newBuilder();
    taskBuilder.getAssignedTask().getTask()
            .setResources(ImmutableSet.of(
                    numCpus(100),
                    ramMb(1024),
                    diskMb(1024)));
    IScheduledTask newTask = IScheduledTask.build(taskBuilder);

    IScheduledTask oldTask1 = buildTaskForJobUpdate(1);
    IScheduledTask oldTask2 = buildTaskForJobUpdate(2);
    storageUtil.expectTaskFetch(
        Query.unscoped().byJob(JOB_KEY).active(),
        newTask, oldTask1, oldTask2);
    expectJobUpdateQuotaCheck(ENOUGH_QUOTA);
    jobUpdateController.start(EasyMock.<IJobUpdate>anyObject(), eq(AUDIT));

    ITaskConfig newConfig = newTask.getAssignedTask().getTask();
    JobUpdate builder = buildJobUpdate(
        3,
        newConfig,
        ImmutableMap.of(
            newConfig, ImmutableSet.of(new Range(0, 0)),
            oldTask1.getAssignedTask().getTask(), ImmutableSet.of(new Range(1, 2))))
        .newBuilder();
    builder.getInstructions().getSettings()
        .setUpdateOnlyTheseInstances(ImmutableSet.of(new Range(0, 2)));

    control.replay();
    JobUpdateRequest request = buildJobUpdateRequest(IJobUpdate.build(builder));
    assertResponse(OK, thrift.startJobUpdate(request, AUDIT_MESSAGE));
    assertEquals(3L, statsProvider.getLongValue(START_JOB_UPDATE));
  }

  @Test
  public void testStartUpdateFailsInstanceCountCheck() throws Exception {
    JobUpdateRequest request = buildServiceJobUpdateRequest(populatedTask());
    request.setInstanceCount(THRESHOLDS.getMaxTasksPerJob() + 1);
    expectGetRemoteUser();
    expectNoCronJob();
    expect(uuidGenerator.createNew()).andReturn(UU_ID);
    storageUtil.expectTaskFetch(Query.unscoped().byJob(JOB_KEY).active());
    expectJobUpdateQuotaCheck(ENOUGH_QUOTA);

    control.replay();

    assertResponse(INVALID_REQUEST, thrift.startJobUpdate(request, AUDIT_MESSAGE));
    assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE));
  }

  @Test
  public void testStartUpdateFailsQuotaCheck() throws Exception {
    JobUpdateRequest request = buildServiceJobUpdateRequest(populatedTask());
    expectGetRemoteUser();
    expectNoCronJob();
    expect(uuidGenerator.createNew()).andReturn(UU_ID);

    IScheduledTask oldTask = buildTaskForJobUpdate(0);
    storageUtil.expectTaskFetch(Query.unscoped().byJob(JOB_KEY).active(), oldTask);

    expectJobUpdateQuotaCheck(NOT_ENOUGH_QUOTA);

    control.replay();

    assertResponse(INVALID_REQUEST, thrift.startJobUpdate(request, AUDIT_MESSAGE));
    assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE));
  }

  @Test
  public void testStartUpdateFailsInController() throws Exception {
    expectGetRemoteUser();
    expectNoCronJob();

    IScheduledTask oldTask = buildTaskForJobUpdate(0, "old");
    ITaskConfig newTask = buildTaskForJobUpdate(0, "new").getAssignedTask().getTask();

    IJobUpdate update = buildJobUpdate(
        1,
        newTask,
        ImmutableMap.of(oldTask.getAssignedTask().getTask(), ImmutableSet.of(new Range(0, 0))));

    expect(uuidGenerator.createNew()).andReturn(UU_ID);
    expectJobUpdateQuotaCheck(ENOUGH_QUOTA);

    storageUtil.expectTaskFetch(Query.unscoped().byJob(JOB_KEY).active(), oldTask);
    jobUpdateController.start(update, AUDIT);
    expectLastCall().andThrow(new UpdateStateException("failed"));

    control.replay();

    assertResponse(
        INVALID_REQUEST,
        thrift.startJobUpdate(buildJobUpdateRequest(update), AUDIT_MESSAGE));
    assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE));
  }

  @Test
  public void testStartUpdateFailsInControllerWhenUpdateInProgress() throws Exception {
    expectGetRemoteUser();
    expectNoCronJob();

    IScheduledTask oldTask = buildTaskForJobUpdate(0, "old");
    ITaskConfig newTask = buildTaskForJobUpdate(0, "new").getAssignedTask().getTask();

    IJobUpdate update = buildJobUpdate(
        1,
        newTask,
        ImmutableMap.of(oldTask.getAssignedTask().getTask(), ImmutableSet.of(new Range(0, 0))));

    expect(uuidGenerator.createNew()).andReturn(UU_ID);
    expectJobUpdateQuotaCheck(ENOUGH_QUOTA);

    storageUtil.expectTaskFetch(Query.unscoped().byJob(JOB_KEY).active(), oldTask);
    jobUpdateController.start(update, AUDIT);
    expectLastCall().andThrow(new UpdateInProgressException("failed", update.getSummary()));

    control.replay();

    Response response = thrift.startJobUpdate(buildJobUpdateRequest(update), AUDIT_MESSAGE);

    assertResponse(INVALID_REQUEST, response);
    assertEquals(
        new StartJobUpdateResult(UPDATE_KEY.newBuilder()).setUpdateSummary(
            update.getSummary().newBuilder()), response.getResult().getStartJobUpdateResult());
    assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE));
  }

  @Test
  public void testStartUpdateFailsWhenSlaAwareWithNoPolicy() throws Exception {
    IScheduledTask oldTask = buildTaskForJobUpdate(0, "old");
    ITaskConfig newTask = buildTaskForJobUpdate(0, "new", Optional.empty())
        .getAssignedTask()
        .getTask();

    IJobUpdate update = buildJobUpdate(
        1,
        newTask,
        ImmutableMap.of(oldTask.getAssignedTask().getTask(), ImmutableSet.of(new Range(0, 0))),
        buildJobUpdateSettings().setSlaAware(true));

    control.replay();

    JobUpdateRequest request = buildJobUpdateRequest(update);
    Response response = thrift.startJobUpdate(request, AUDIT_MESSAGE);
    assertEquals(invalidResponse(INVALID_SLA_AWARE_UPDATE), response);
    assertEquals(0L, statsProvider.getLongValue(START_JOB_UPDATE));
  }

  @Test
  public void testPauseJobUpdateByCoordinator() throws Exception {
    expectGetRemoteUser();
    jobUpdateController.pause(UPDATE_KEY, AUDIT);

    control.replay();

    assertResponse(OK, thrift.pauseJobUpdate(UPDATE_KEY.newBuilder(), AUDIT_MESSAGE));
  }

  @Test
  public void testPauseJobUpdateByUser() throws Exception {
    expectGetRemoteUser();
    jobUpdateController.pause(UPDATE_KEY, AUDIT);

    control.replay();

    assertResponse(OK, thrift.pauseJobUpdate(UPDATE_KEY.newBuilder(), AUDIT_MESSAGE));
  }

  @Test(expected = IllegalArgumentException.class)
  public void testPauseMessageTooLong() throws Exception {
    expectGetRemoteUser();

    control.replay();

    assertResponse(
        OK,
        thrift.pauseJobUpdate(
            UPDATE_KEY.newBuilder(),
            Strings.repeat("*", AuditData.MAX_MESSAGE_LENGTH + 1)));
  }

  @Test
  public void testPauseJobUpdateFailsInController() throws Exception {
    expectGetRemoteUser();
    jobUpdateController.pause(UPDATE_KEY, AUDIT);
    expectLastCall().andThrow(new UpdateStateException("failed"));

    control.replay();

    assertResponse(
        INVALID_REQUEST,
        thrift.pauseJobUpdate(UPDATE_KEY.newBuilder(), AUDIT_MESSAGE));
  }

  @Test
  public void testResumeJobUpdate() throws Exception {
    expectGetRemoteUser();
    jobUpdateController.resume(UPDATE_KEY, AUDIT);

    control.replay();

    assertResponse(OK, thrift.resumeJobUpdate(UPDATE_KEY.newBuilder(), AUDIT_MESSAGE));
  }

  @Test
  public void testResumeJobUpdateFailsInController() throws Exception {
    expectGetRemoteUser();
    jobUpdateController.resume(UPDATE_KEY, AUDIT);
    expectLastCall().andThrow(new UpdateStateException("failed"));

    control.replay();

    assertResponse(
        INVALID_REQUEST,
        thrift.resumeJobUpdate(UPDATE_KEY.newBuilder(), AUDIT_MESSAGE));
  }

  @Test
  public void testAbortJobUpdateByCoordinator() throws Exception {
    expectGetRemoteUser();
    jobUpdateController.abort(UPDATE_KEY, AUDIT);

    control.replay();

    assertResponse(OK, thrift.abortJobUpdate(UPDATE_KEY.newBuilder(), AUDIT_MESSAGE));
  }

  @Test
  public void testAbortJobUpdateByUser() throws Exception {
    expectGetRemoteUser();
    jobUpdateController.abort(UPDATE_KEY, AUDIT);

    control.replay();

    assertResponse(OK, thrift.abortJobUpdate(UPDATE_KEY.newBuilder(), AUDIT_MESSAGE));
  }

  @Test
  public void testAbortJobUpdateFailsInController() throws Exception {
    expectGetRemoteUser();
    jobUpdateController.abort(UPDATE_KEY, AUDIT);
    expectLastCall().andThrow(new UpdateStateException("failed"));

    control.replay();

    assertResponse(
        INVALID_REQUEST,
        thrift.abortJobUpdate(UPDATE_KEY.newBuilder(), AUDIT_MESSAGE));
  }

  @Test
  public void testPulseJobUpdatePulsedAsCoordinator() throws Exception {
    expect(jobUpdateController.pulse(UPDATE_KEY)).andReturn(JobUpdatePulseStatus.OK);

    control.replay();

    assertEquals(
        okResponse(Result.pulseJobUpdateResult(
            new PulseJobUpdateResult(JobUpdatePulseStatus.OK))),
        thrift.pulseJobUpdate(UPDATE_KEY.newBuilder()));
  }

  @Test
  public void testPulseJobUpdatePulsedAsUser() throws Exception {
    expect(jobUpdateController.pulse(UPDATE_KEY)).andReturn(JobUpdatePulseStatus.OK);

    control.replay();

    assertEquals(
        okResponse(Result.pulseJobUpdateResult(new PulseJobUpdateResult(JobUpdatePulseStatus.OK))),
        thrift.pulseJobUpdate(UPDATE_KEY.newBuilder()));
  }

  @Test
  public void testPulseJobUpdateFails() throws Exception {
    expect(jobUpdateController.pulse(UPDATE_KEY)).andThrow(new UpdateStateException("failure"));

    control.replay();

    assertResponse(INVALID_REQUEST, thrift.pulseJobUpdate(UPDATE_KEY.newBuilder()));
  }

  @Test
  public void testGetJobUpdateSummaries() throws Exception {
    Response updateSummary = Responses.empty()
        .setResponseCode(OK)
        .setDetails(ImmutableList.of(new ResponseDetail("summary")));

    expect(readOnlyScheduler.getJobUpdateSummaries(
        anyObject(JobUpdateQuery.class))).andReturn(updateSummary);
    control.replay();

    assertEquals(updateSummary, thrift.getJobUpdateSummaries(new JobUpdateQuery()));
  }

  private IExpectationSetters<String> expectGetRemoteUser() {
    return expect(auditMessages.getRemoteUserName()).andReturn(USER);
  }

  private static TaskConfig populatedTask() {
    return defaultTask(true).setConstraints(Sets.newHashSet(
        new Constraint("host", TaskConstraint.limit(new LimitConstraint(1)))));
  }

  private static Constraint dedicatedConstraint(int value) {
    return new Constraint(DEDICATED_ATTRIBUTE, TaskConstraint.limit(new LimitConstraint(value)));
  }

  private static Constraint dedicatedConstraint(Set<String> values) {
    return new Constraint(DEDICATED_ATTRIBUTE,
        TaskConstraint.value(new ValueConstraint(false, values)));
  }

  private static JobUpdateRequest buildServiceJobUpdateRequest() {
    return buildServiceJobUpdateRequest(defaultTask(true));
  }

  private static JobUpdateRequest buildServiceJobUpdateRequest(TaskConfig config) {
    return buildJobUpdateRequest(config.setIsService(true));
  }

  private static JobUpdateRequest buildJobUpdateRequest(TaskConfig config) {
    return new JobUpdateRequest()
        .setInstanceCount(6)
        .setSettings(buildJobUpdateSettings())
        .setTaskConfig(config);
  }

  private static JobUpdateSettings buildJobUpdateSettings() {
    return new JobUpdateSettings()
        .setUpdateGroupSize(10)
        .setUpdateStrategy(
            JobUpdateStrategy.queueStrategy(new QueueJobUpdateStrategy().setGroupSize(10)))
        .setMaxFailedInstances(2)
        .setMaxPerInstanceFailures(1)
        .setMinWaitInInstanceRunningMs(15000)
        .setRollbackOnFailure(true);
  }

  private static Integer rangesToInstanceCount(Set<IRange> ranges) {
    int instanceCount = 0;
    for (IRange range : ranges) {
      instanceCount += range.getLast() - range.getFirst() + 1;
    }

    return instanceCount;
  }

  private static JobUpdateRequest buildJobUpdateRequest(IJobUpdate update) {
    return new JobUpdateRequest()
        .setInstanceCount(rangesToInstanceCount(
            update.getInstructions().getDesiredState().getInstances()))
        .setSettings(update.getInstructions().getSettings().newBuilder())
        .setTaskConfig(update.getInstructions().getDesiredState().getTask().newBuilder())
        .setMetadata(IMetadata.toBuildersSet(update.getSummary().getMetadata()));
  }

  private static IJobUpdate buildJobUpdate(
      int instanceCount,
      ITaskConfig newConfig,
      ImmutableMap<ITaskConfig, ImmutableSet<Range>> oldConfigMap) {

    return buildJobUpdate(instanceCount, newConfig, oldConfigMap, buildJobUpdateSettings());
  }

  private static IJobUpdate buildJobUpdate(
      int instanceCount,
      ITaskConfig newConfig,
      ImmutableMap<ITaskConfig, ImmutableSet<Range>> oldConfigMap,
      JobUpdateSettings jobUpdateSettings) {

    ImmutableSet.Builder<InstanceTaskConfig> builder = ImmutableSet.builder();
    for (Map.Entry<ITaskConfig, ImmutableSet<Range>> entry : oldConfigMap.entrySet()) {
      builder.add(new InstanceTaskConfig(entry.getKey().newBuilder(), entry.getValue()));
    }

    return IJobUpdate.build(new JobUpdate()
        .setSummary(new JobUpdateSummary()
            .setKey(UPDATE_KEY.newBuilder())
            .setUser(IDENTITY.getUser())
            .setMetadata(METADATA))
        .setInstructions(new JobUpdateInstructions()
            .setSettings(jobUpdateSettings)
            .setDesiredState(new InstanceTaskConfig()
                .setTask(newConfig.newBuilder())
                .setInstances(ImmutableSet.of(new Range(0, instanceCount - 1))))
            .setInitialState(builder.build())));
  }

  private IExpectationSetters<?> expectCronJob() {
    return expect(storageUtil.jobStore.fetchJob(JOB_KEY))
        .andReturn(Optional.of(IJobConfiguration.build(CRON_JOB)));
  }

  private IExpectationSetters<?> expectNoCronJob() {
    return expect(storageUtil.jobStore.fetchJob(JOB_KEY))
        .andReturn(Optional.empty());
  }

  private IExpectationSetters<?> expectInstanceQuotaCheck(
      SanitizedConfiguration sanitized,
      QuotaCheckResult result) {

    return expectInstanceQuotaCheck(sanitized.getJobConfig().getTaskConfig(), result);
  }

  private IExpectationSetters<?> expectInstanceQuotaCheck(
      ITaskConfig config,
      QuotaCheckResult result) {

    return expect(quotaManager.checkInstanceAddition(
        config,
        1,
        storageUtil.mutableStoreProvider)).andReturn(result);
  }

  private IExpectationSetters<?> expectCronQuotaCheck(
      IJobConfiguration config,
      QuotaCheckResult result) {

    return expect(quotaManager.checkCronUpdate(config, storageUtil.mutableStoreProvider))
        .andReturn(result);
  }
}