@Test
  public void testGetTasksStatusPagination() throws Exception {
    Iterable<IScheduledTask> tasks = makeDefaultScheduledTasks(10);

    TaskQuery page1Query = setupPaginatedQuery(tasks, 0, 4);
    TaskQuery page2Query = setupPaginatedQuery(tasks, 4, 4);
    TaskQuery page3Query = setupPaginatedQuery(tasks, 8, 4);

    control.replay();

    Response page1Response = assertOkResponse(thrift.getTasksStatus(page1Query));
    Response page2Response = assertOkResponse(thrift.getTasksStatus(page2Query));
    Response page3Response = assertOkResponse(thrift.getTasksStatus(page3Query));

    Iterable<Integer> page1Ids =
        Lists.newArrayList(
            Iterables.transform(
                page1Response.getResult().getScheduleStatusResult().getTasks(), TO_INSTANCE_ID));
    Iterable<Integer> page2Ids =
        Lists.newArrayList(
            Iterables.transform(
                page2Response.getResult().getScheduleStatusResult().getTasks(), TO_INSTANCE_ID));
    Iterable<Integer> page3Ids =
        Lists.newArrayList(
            Iterables.transform(
                page3Response.getResult().getScheduleStatusResult().getTasks(), TO_INSTANCE_ID));

    assertEquals(Lists.newArrayList(0, 1, 2, 3), page1Ids);
    assertEquals(Lists.newArrayList(4, 5, 6, 7), page2Ids);
    assertEquals(Lists.newArrayList(8, 9), page3Ids);
  }
  @Test
  public void testPopulateJobConfigFails() throws Exception {
    IJobConfiguration job = IJobConfiguration.build(makeJob(null));
    control.replay();

    assertResponse(INVALID_REQUEST, thrift.populateJobConfig(job.newBuilder()));
  }
  @Test
  public void testGetJobUpdateDiffWithUpdateRemove() throws Exception {
    TaskConfig task1 = defaultTask(false).setNumCpus(1.0);
    TaskConfig task2 = defaultTask(false).setNumCpus(2.0);
    TaskConfig task3 = defaultTask(false).setNumCpus(3.0);

    ImmutableSet.Builder<IScheduledTask> tasks = ImmutableSet.builder();
    makeTasks(0, 10, task1, tasks);
    makeTasks(10, 20, task2, tasks);
    makeTasks(20, 30, task3, tasks);

    expect(storageUtil.jobStore.fetchJob(JOB_KEY)).andReturn(Optional.absent());
    storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active(), tasks.build());

    control.replay();

    JobUpdateRequest request =
        new JobUpdateRequest()
            .setTaskConfig(defaultTask(false).setNumCpus(6.0))
            .setInstanceCount(20)
            .setSettings(new JobUpdateSettings());

    GetJobUpdateDiffResult expected =
        new GetJobUpdateDiffResult()
            .setRemove(ImmutableSet.of(group(task3, new Range(20, 29))))
            .setUpdate(
                ImmutableSet.of(group(task1, new Range(0, 9)), group(task2, new Range(10, 19))))
            .setAdd(ImmutableSet.of())
            .setUnchanged(ImmutableSet.of());

    Response response = assertOkResponse(thrift.getJobUpdateDiff(request));
    assertEquals(expected, response.getResult().getGetJobUpdateDiffResult());
  }
  @Test
  public void testGetPendingReasonFailsStatusSet() throws Exception {
    Builder query = Query.unscoped().byStatus(ScheduleStatus.ASSIGNED);

    control.replay();

    assertResponse(INVALID_REQUEST, thrift.getPendingReason(query.get()));
  }
  @Test
  public void testGetPendingReasonFailsSlavesSet() throws Exception {
    Builder query = Query.unscoped().bySlave("host1");

    control.replay();

    assertResponse(INVALID_REQUEST, thrift.getPendingReason(query.get()));
  }
  @Test
  public void testGetJobUpdateDiffInvalidConfig() throws Exception {
    control.replay();

    JobUpdateRequest request =
        new JobUpdateRequest().setTaskConfig(defaultTask(false).setNumCpus(-1));
    assertResponse(INVALID_REQUEST, thrift.getJobUpdateDiff(request));
  }
  @Test
  public void testGetJobUpdateDetailsInvalidId() throws Exception {
    expect(storageUtil.jobUpdateStore.fetchJobUpdateDetails(UPDATE_KEY))
        .andReturn(Optional.absent());

    control.replay();

    assertResponse(INVALID_REQUEST, thrift.getJobUpdateDetails(UPDATE_KEY.newBuilder()));
  }
  @Test
  public void testGetLocks() throws Exception {
    expect(lockManager.getLocks()).andReturn(ImmutableSet.of(LOCK));

    control.replay();

    Response response = thrift.getLocks();
    assertEquals(
        LOCK.newBuilder(),
        Iterables.getOnlyElement(response.getResult().getGetLocksResult().getLocks()));
  }
  @Test
  public void testGetTasksStatus() throws Exception {
    Builder query = Query.unscoped();
    Iterable<IScheduledTask> tasks = makeDefaultScheduledTasks(10);
    storageUtil.expectTaskFetch(query, ImmutableSet.copyOf(tasks));

    control.replay();

    ImmutableList<ScheduledTask> expected = IScheduledTask.toBuildersList(tasks);
    Response response = assertOkResponse(thrift.getTasksStatus(new TaskQuery()));
    assertEquals(expected, response.getResult().getScheduleStatusResult().getTasks());
  }
  @Test
  public void testGetJobUpdateDetails() throws Exception {
    JobUpdateDetails details = createJobUpdateDetails();
    expect(storageUtil.jobUpdateStore.fetchJobUpdateDetails(UPDATE_KEY))
        .andReturn(Optional.of(IJobUpdateDetails.build(details)));

    control.replay();

    Response response = assertOkResponse(thrift.getJobUpdateDetails(UPDATE_KEY.newBuilder()));
    assertEquals(
        IJobUpdateDetails.build(details),
        IJobUpdateDetails.build(response.getResult().getGetJobUpdateDetailsResult().getDetails()));
  }
  @Test
  public void testPopulateJobConfig() throws Exception {
    IJobConfiguration job = IJobConfiguration.build(makeJob());
    SanitizedConfiguration sanitized =
        SanitizedConfiguration.fromUnsanitized(TaskTestUtil.CONFIGURATION_MANAGER, job);
    control.replay();

    Response response = assertOkResponse(thrift.populateJobConfig(job.newBuilder()));
    assertEquals(
        Result.populateJobResult(
            new PopulateJobResult(sanitized.getJobConfig().getTaskConfig().newBuilder())),
        response.getResult());
  }
  @Test
  public void testGetJobUpdateSummaries() throws Exception {
    JobUpdateQuery query = new JobUpdateQuery().setRole(ROLE);
    List<JobUpdateSummary> summaries = createJobUpdateSummaries(5);
    expect(storageUtil.jobUpdateStore.fetchJobUpdateSummaries(IJobUpdateQuery.build(query)))
        .andReturn(IJobUpdateSummary.listFromBuilders(summaries));

    control.replay();

    Response response = assertOkResponse(thrift.getJobUpdateSummaries(query));
    assertEquals(
        summaries, response.getResult().getGetJobUpdateSummariesResult().getUpdateSummaries());
  }
  @Test
  public void testGetTasksWithoutConfigs() throws Exception {
    Builder query = Query.unscoped();
    storageUtil.expectTaskFetch(query, ImmutableSet.copyOf(makeDefaultScheduledTasks(10)));

    control.replay();

    ImmutableList<ScheduledTask> expected =
        IScheduledTask.toBuildersList(
            makeDefaultScheduledTasks(10, defaultTask(true).setExecutorConfig(null)));

    Response response = assertOkResponse(thrift.getTasksWithoutConfigs(new TaskQuery()));
    assertEquals(expected, response.getResult().getScheduleStatusResult().getTasks());
  }
  @Test
  public void testEmptyConfigSummary() throws Exception {
    IJobKey key = JobKeys.from("test", "test", "test");

    storageUtil.expectTaskFetch(Query.jobScoped(key).active(), ImmutableSet.of());

    ConfigSummary summary =
        new ConfigSummary().setKey(key.newBuilder()).setGroups(Sets.newHashSet());

    ConfigSummaryResult expected = new ConfigSummaryResult().setSummary(summary);

    control.replay();

    Response response = assertOkResponse(thrift.getConfigSummary(key.newBuilder()));
    assertEquals(expected, response.getResult().getConfigSummaryResult());
  }
  @Test
  public void testGetJobUpdateDiffNoCron() throws Exception {
    expect(storageUtil.jobStore.fetchJob(JOB_KEY))
        .andReturn(Optional.of(IJobConfiguration.build(CRON_JOB)));

    control.replay();

    JobUpdateRequest request = new JobUpdateRequest().setTaskConfig(defaultTask(false));

    Response expected =
        Responses.empty()
            .setResponseCode(INVALID_REQUEST)
            .setDetails(ImmutableList.of(new ResponseDetail(NO_CRON)));

    assertEquals(expected, thrift.getJobUpdateDiff(request));
  }
  @Test
  public void testGetConfigSummary() throws Exception {
    IJobKey key = JobKeys.from("test", "test", "test");

    TaskConfig firstGroupTask = defaultTask(true);
    TaskConfig secondGroupTask = defaultTask(true).setNumCpus(2);

    IScheduledTask first1 =
        IScheduledTask.build(
            new ScheduledTask()
                .setAssignedTask(new AssignedTask().setTask(firstGroupTask).setInstanceId(0)));

    IScheduledTask first2 =
        IScheduledTask.build(
            new ScheduledTask()
                .setAssignedTask(new AssignedTask().setTask(firstGroupTask).setInstanceId(1)));

    IScheduledTask second =
        IScheduledTask.build(
            new ScheduledTask()
                .setAssignedTask(new AssignedTask().setTask(secondGroupTask).setInstanceId(2)));

    storageUtil.expectTaskFetch(Query.jobScoped(key).active(), first1, first2, second);

    ConfigGroup group1 =
        new ConfigGroup()
            .setConfig(firstGroupTask)
            .setInstances(IRange.toBuildersSet(convertRanges(toRanges(ImmutableSet.of(0, 1)))));
    ConfigGroup group2 =
        new ConfigGroup()
            .setConfig(secondGroupTask)
            .setInstances(IRange.toBuildersSet(convertRanges(toRanges(ImmutableSet.of(2)))));

    ConfigSummary summary =
        new ConfigSummary().setKey(key.newBuilder()).setGroups(Sets.newHashSet(group1, group2));

    ConfigSummaryResult expected = new ConfigSummaryResult().setSummary(summary);

    control.replay();

    Response response = assertOkResponse(thrift.getConfigSummary(key.newBuilder()));
    assertEquals(
        IConfigSummaryResult.build(expected),
        IConfigSummaryResult.build(response.getResult().getConfigSummaryResult()));
  }
  @Test
  public void testGetAllJobs() throws Exception {
    JobConfiguration cronJobOne =
        makeJob()
            .setCronSchedule("1 * * * *")
            .setKey(JOB_KEY.newBuilder())
            .setTaskConfig(nonProductionTask());
    JobKey jobKey2 = JOB_KEY.newBuilder().setRole("other_role");
    JobConfiguration cronJobTwo =
        makeJob().setCronSchedule("2 * * * *").setKey(jobKey2).setTaskConfig(nonProductionTask());
    TaskConfig immediateTaskConfig =
        defaultTask(false)
            .setJob(JOB_KEY.newBuilder().setName("immediate"))
            .setJobName("immediate")
            .setOwner(ROLE_IDENTITY);
    IScheduledTask immediateTask =
        IScheduledTask.build(
            new ScheduledTask()
                .setAssignedTask(new AssignedTask().setTask(immediateTaskConfig))
                .setStatus(ScheduleStatus.ASSIGNED));
    JobConfiguration immediateJob =
        new JobConfiguration()
            .setKey(JOB_KEY.newBuilder().setName("immediate"))
            .setOwner(ROLE_IDENTITY)
            .setInstanceCount(1)
            .setTaskConfig(immediateTaskConfig);

    Set<JobConfiguration> crons = ImmutableSet.of(cronJobOne, cronJobTwo);
    expect(storageUtil.jobStore.fetchJobs()).andReturn(IJobConfiguration.setFromBuilders(crons));
    storageUtil.expectTaskFetch(Query.unscoped().active(), immediateTask);

    control.replay();

    Set<JobConfiguration> allJobs =
        ImmutableSet.<JobConfiguration>builder().addAll(crons).add(immediateJob).build();
    assertEquals(
        IJobConfiguration.setFromBuilders(allJobs),
        IJobConfiguration.setFromBuilders(
            thrift.getJobs(null).getResult().getGetJobsResult().getConfigs()));
  }
  @Test
  public void testGetQuota() throws Exception {
    QuotaInfo infoMock = createMock(QuotaInfo.class);
    expect(quotaManager.getQuotaInfo(ROLE, storageUtil.storeProvider)).andReturn(infoMock);
    expect(infoMock.getQuota()).andReturn(QUOTA);
    expect(infoMock.getProdSharedConsumption()).andReturn(XLARGE);
    expect(infoMock.getProdDedicatedConsumption()).andReturn(LARGE);
    expect(infoMock.getNonProdSharedConsumption()).andReturn(MEDIUM);
    expect(infoMock.getNonProdDedicatedConsumption()).andReturn(SMALL);
    control.replay();

    GetQuotaResult expected =
        new GetQuotaResult()
            .setQuota(QUOTA.newBuilder())
            .setProdSharedConsumption(XLARGE.newBuilder())
            .setProdDedicatedConsumption(LARGE.newBuilder())
            .setNonProdSharedConsumption(MEDIUM.newBuilder())
            .setNonProdDedicatedConsumption(SMALL.newBuilder());

    Response response = assertOkResponse(thrift.getQuota(ROLE));
    assertEquals(expected, response.getResult().getGetQuotaResult());
  }
  @Test
  public void testGetPendingReason() throws Exception {
    Builder query = Query.unscoped().byJob(JOB_KEY);
    Builder filterQuery = Query.unscoped().byJob(JOB_KEY).byStatus(ScheduleStatus.PENDING);
    String taskId1 = "task_id_test1";
    String taskId2 = "task_id_test2";
    ImmutableSet<Veto> result =
        ImmutableSet.of(Veto.constraintMismatch("first"), Veto.constraintMismatch("second"));

    ITaskConfig taskConfig = ITaskConfig.build(defaultTask(true));
    IScheduledTask pendingTask1 =
        IScheduledTask.build(
            new ScheduledTask()
                .setAssignedTask(
                    new AssignedTask().setTaskId(taskId1).setTask(taskConfig.newBuilder()))
                .setStatus(ScheduleStatus.PENDING));

    IScheduledTask pendingTask2 =
        IScheduledTask.build(
            new ScheduledTask()
                .setAssignedTask(
                    new AssignedTask().setTaskId(taskId2).setTask(taskConfig.newBuilder()))
                .setStatus(ScheduleStatus.PENDING));

    storageUtil.expectTaskFetch(filterQuery, pendingTask1, pendingTask2);
    expect(nearestFit.getNearestFit(TaskGroupKey.from(taskConfig))).andReturn(result).times(2);

    control.replay();

    String reason = "Constraint not satisfied: first,Constraint not satisfied: second";
    Set<PendingReason> expected =
        ImmutableSet.of(
            new PendingReason().setTaskId(taskId1).setReason(reason),
            new PendingReason().setTaskId(taskId2).setReason(reason));

    Response response = assertOkResponse(thrift.getPendingReason(query.get()));
    assertEquals(expected, response.getResult().getGetPendingReasonResult().getReasons());
  }
  @Test
  public void testGetJobUpdateDiffWithUnchanged() throws Exception {
    expect(storageUtil.jobStore.fetchJob(JOB_KEY)).andReturn(Optional.absent());
    storageUtil.expectTaskFetch(
        Query.jobScoped(JOB_KEY).active(), ImmutableSet.copyOf(makeDefaultScheduledTasks(10)));

    control.replay();

    JobUpdateRequest request =
        new JobUpdateRequest()
            .setTaskConfig(defaultTask(true))
            .setInstanceCount(10)
            .setSettings(new JobUpdateSettings());

    GetJobUpdateDiffResult expected =
        new GetJobUpdateDiffResult()
            .setUnchanged(ImmutableSet.of(group(defaultTask(true), new Range(0, 9))))
            .setRemove(ImmutableSet.of())
            .setUpdate(ImmutableSet.of())
            .setAdd(ImmutableSet.of());

    Response response = assertOkResponse(thrift.getJobUpdateDiff(request));
    assertEquals(expected, response.getResult().getGetJobUpdateDiffResult());
  }
  @Test
  public void testGetJobSummaryWithoutNextRun() throws Exception {
    // 31st of February, there is no such day.
    String cronSchedule = "* * 31 2 *";

    TaskConfig task =
        nonProductionTask()
            .setJobName(JOB_KEY.getName())
            .setOwner(ROLE_IDENTITY)
            .setEnvironment(JOB_KEY.getEnvironment());
    JobConfiguration job = makeJob().setCronSchedule(cronSchedule).setTaskConfig(task);
    expect(cronPredictor.predictNextRun(CrontabEntry.parse(cronSchedule)))
        .andReturn(Optional.absent())
        .anyTimes();
    storageUtil.expectTaskFetch(Query.roleScoped(ROLE));
    Set<JobConfiguration> jobOnly = ImmutableSet.of(job);
    expect(storageUtil.jobStore.fetchJobs()).andReturn(IJobConfiguration.setFromBuilders(jobOnly));

    control.replay();

    JobSummaryResult result = thrift.getJobSummary(ROLE).getResult().getJobSummaryResult();
    assertEquals(1, result.getSummaries().size());
    assertFalse(result.getSummariesIterator().next().isSetNextCronRunMs());
  }
  @Test
  public void testGetJobSummary() throws Exception {
    long nextCronRunMs = 100;
    TaskConfig ownedCronJobTask =
        nonProductionTask()
            .setJob(JOB_KEY.newBuilder())
            .setJobName(JOB_KEY.getName())
            .setOwner(ROLE_IDENTITY)
            .setEnvironment(JOB_KEY.getEnvironment());
    JobConfiguration ownedCronJob =
        makeJob().setCronSchedule(CRON_SCHEDULE).setTaskConfig(ownedCronJobTask);
    IScheduledTask ownedCronJobScheduledTask =
        IScheduledTask.build(
            new ScheduledTask()
                .setAssignedTask(new AssignedTask().setTask(ownedCronJobTask))
                .setStatus(ScheduleStatus.ASSIGNED));
    Identity otherOwner = new Identity("other", "other");
    JobConfiguration unownedCronJob =
        makeJob()
            .setOwner(otherOwner)
            .setCronSchedule(CRON_SCHEDULE)
            .setKey(JOB_KEY.newBuilder().setRole("other"))
            .setTaskConfig(ownedCronJobTask.deepCopy().setOwner(otherOwner));
    TaskConfig ownedImmediateTaskInfo =
        defaultTask(false)
            .setJob(JOB_KEY.newBuilder().setName("immediate"))
            .setJobName("immediate")
            .setOwner(ROLE_IDENTITY);
    Set<JobConfiguration> ownedCronJobOnly = ImmutableSet.of(ownedCronJob);
    Set<JobSummary> ownedCronJobSummaryOnly =
        ImmutableSet.of(
            new JobSummary()
                .setJob(ownedCronJob)
                .setStats(new JobStats())
                .setNextCronRunMs(nextCronRunMs));
    Set<JobSummary> ownedCronJobSummaryWithRunningTask =
        ImmutableSet.of(
            new JobSummary()
                .setJob(ownedCronJob)
                .setStats(new JobStats().setActiveTaskCount(1))
                .setNextCronRunMs(nextCronRunMs));
    Set<JobConfiguration> unownedCronJobOnly = ImmutableSet.of(unownedCronJob);
    Set<JobConfiguration> bothCronJobs = ImmutableSet.of(ownedCronJob, unownedCronJob);

    IScheduledTask ownedImmediateTask =
        IScheduledTask.build(
            new ScheduledTask()
                .setAssignedTask(new AssignedTask().setTask(ownedImmediateTaskInfo))
                .setStatus(ScheduleStatus.ASSIGNED));
    JobConfiguration ownedImmediateJob =
        new JobConfiguration()
            .setKey(JOB_KEY.newBuilder().setName("immediate"))
            .setOwner(ROLE_IDENTITY)
            .setInstanceCount(1)
            .setTaskConfig(ownedImmediateTaskInfo);
    Builder query = Query.roleScoped(ROLE);

    Set<JobSummary> ownedImmediateJobSummaryOnly =
        ImmutableSet.of(
            new JobSummary()
                .setJob(ownedImmediateJob)
                .setStats(new JobStats().setActiveTaskCount(1)));

    expect(cronPredictor.predictNextRun(CrontabEntry.parse(CRON_SCHEDULE)))
        .andReturn(Optional.of(new Date(nextCronRunMs)))
        .anyTimes();

    storageUtil.expectTaskFetch(query);
    expect(storageUtil.jobStore.fetchJobs())
        .andReturn(IJobConfiguration.setFromBuilders(ownedCronJobOnly));

    storageUtil.expectTaskFetch(query);
    expect(storageUtil.jobStore.fetchJobs())
        .andReturn(IJobConfiguration.setFromBuilders(bothCronJobs));

    storageUtil.expectTaskFetch(query, ownedImmediateTask);
    expect(storageUtil.jobStore.fetchJobs())
        .andReturn(IJobConfiguration.setFromBuilders(unownedCronJobOnly));

    storageUtil.expectTaskFetch(query);
    expect(storageUtil.jobStore.fetchJobs()).andReturn(ImmutableSet.of());

    // Handle the case where a cron job has a running task (same JobKey present in both stores).
    storageUtil.expectTaskFetch(query, ownedCronJobScheduledTask);
    expect(storageUtil.jobStore.fetchJobs())
        .andReturn(IJobConfiguration.setFromBuilders(ImmutableSet.of(ownedCronJob)));

    control.replay();

    assertEquals(jobSummaryResponse(ownedCronJobSummaryOnly), thrift.getJobSummary(ROLE));

    assertEquals(jobSummaryResponse(ownedCronJobSummaryOnly), thrift.getJobSummary(ROLE));

    Response jobSummaryResponse = thrift.getJobSummary(ROLE);
    assertEquals(
        jobSummaryResponse(ownedImmediateJobSummaryOnly),
        IResponse.build(jobSummaryResponse).newBuilder());

    assertEquals(jobSummaryResponse(ImmutableSet.of()), thrift.getJobSummary(ROLE));

    assertEquals(
        jobSummaryResponse(ownedCronJobSummaryWithRunningTask), thrift.getJobSummary(ROLE));
  }
  @Test
  public void testGetJobs() throws Exception {
    TaskConfig ownedCronJobTask =
        nonProductionTask()
            .setJobName(JOB_KEY.getName())
            .setOwner(ROLE_IDENTITY)
            .setEnvironment(JOB_KEY.getEnvironment());
    JobConfiguration ownedCronJob =
        makeJob().setCronSchedule(CRON_SCHEDULE).setTaskConfig(ownedCronJobTask);
    IScheduledTask ownedCronJobScheduledTask =
        IScheduledTask.build(
            new ScheduledTask()
                .setAssignedTask(new AssignedTask().setTask(ownedCronJobTask))
                .setStatus(ScheduleStatus.ASSIGNED));
    Identity otherOwner = new Identity("other", "other");
    JobConfiguration unownedCronJob =
        makeJob()
            .setOwner(otherOwner)
            .setCronSchedule(CRON_SCHEDULE)
            .setKey(JOB_KEY.newBuilder().setRole("other"))
            .setTaskConfig(ownedCronJobTask.deepCopy().setOwner(otherOwner));
    TaskConfig ownedImmediateTaskInfo =
        defaultTask(false)
            .setJob(JOB_KEY.newBuilder().setName("immediate"))
            .setJobName("immediate")
            .setOwner(ROLE_IDENTITY);
    Set<JobConfiguration> ownedCronJobOnly = ImmutableSet.of(ownedCronJob);
    Set<JobConfiguration> unownedCronJobOnly = ImmutableSet.of(unownedCronJob);
    Set<JobConfiguration> bothCronJobs = ImmutableSet.of(ownedCronJob, unownedCronJob);
    IScheduledTask ownedImmediateTask =
        IScheduledTask.build(
            new ScheduledTask()
                .setAssignedTask(new AssignedTask().setTask(ownedImmediateTaskInfo))
                .setStatus(ScheduleStatus.ASSIGNED));
    JobConfiguration ownedImmediateJob =
        new JobConfiguration()
            .setKey(JOB_KEY.newBuilder().setName("immediate"))
            .setOwner(ROLE_IDENTITY)
            .setInstanceCount(1)
            .setTaskConfig(ownedImmediateTaskInfo);
    Query.Builder query = Query.roleScoped(ROLE).active();

    storageUtil.expectTaskFetch(query);
    expect(storageUtil.jobStore.fetchJobs())
        .andReturn(IJobConfiguration.setFromBuilders(ownedCronJobOnly));

    storageUtil.expectTaskFetch(query);
    expect(storageUtil.jobStore.fetchJobs())
        .andReturn(IJobConfiguration.setFromBuilders(bothCronJobs));

    storageUtil.expectTaskFetch(query, ownedImmediateTask);
    expect(storageUtil.jobStore.fetchJobs())
        .andReturn(IJobConfiguration.setFromBuilders(unownedCronJobOnly));

    expect(storageUtil.jobStore.fetchJobs()).andReturn(ImmutableSet.of());
    storageUtil.expectTaskFetch(query);

    // Handle the case where a cron job has a running task (same JobKey present in both stores).
    storageUtil.expectTaskFetch(query, ownedCronJobScheduledTask);
    expect(storageUtil.jobStore.fetchJobs())
        .andReturn(IJobConfiguration.setFromBuilders(ImmutableSet.of(ownedCronJob)));

    control.replay();

    assertJobsEqual(
        ownedCronJob,
        Iterables.getOnlyElement(thrift.getJobs(ROLE).getResult().getGetJobsResult().getConfigs()));

    assertJobsEqual(
        ownedCronJob,
        Iterables.getOnlyElement(thrift.getJobs(ROLE).getResult().getGetJobsResult().getConfigs()));

    Set<JobConfiguration> queryResult3 =
        thrift.getJobs(ROLE).getResult().getGetJobsResult().getConfigs();
    assertJobsEqual(ownedImmediateJob, Iterables.getOnlyElement(queryResult3));
    assertEquals(
        ITaskConfig.build(ownedImmediateTaskInfo),
        ITaskConfig.build(Iterables.getOnlyElement(queryResult3).getTaskConfig()));

    assertTrue(thrift.getJobs(ROLE).getResult().getGetJobsResult().getConfigs().isEmpty());

    assertJobsEqual(
        ownedCronJob,
        Iterables.getOnlyElement(thrift.getJobs(ROLE).getResult().getGetJobsResult().getConfigs()));
  }
  @Test
  public void testGetRoleSummary() throws Exception {
    final String BAZ_ROLE = "baz_role";
    final Identity BAZ_ROLE_IDENTITY = new Identity(BAZ_ROLE, USER);

    JobConfiguration cronJobOne =
        makeJob()
            .setCronSchedule("1 * * * *")
            .setKey(JOB_KEY.newBuilder())
            .setTaskConfig(nonProductionTask());
    JobConfiguration cronJobTwo =
        makeJob()
            .setCronSchedule("2 * * * *")
            .setKey(JOB_KEY.newBuilder().setName("cronJob2"))
            .setTaskConfig(nonProductionTask());

    JobConfiguration cronJobThree =
        makeJob()
            .setCronSchedule("3 * * * *")
            .setKey(JOB_KEY.newBuilder().setRole(BAZ_ROLE))
            .setTaskConfig(nonProductionTask())
            .setOwner(BAZ_ROLE_IDENTITY);

    Set<JobConfiguration> crons = ImmutableSet.of(cronJobOne, cronJobTwo, cronJobThree);

    TaskConfig immediateTaskConfig =
        defaultTask(false)
            .setJob(JOB_KEY.newBuilder().setName("immediate"))
            .setJobName("immediate")
            .setOwner(ROLE_IDENTITY);
    IScheduledTask task1 =
        IScheduledTask.build(
            new ScheduledTask().setAssignedTask(new AssignedTask().setTask(immediateTaskConfig)));
    IScheduledTask task2 =
        IScheduledTask.build(
            new ScheduledTask()
                .setAssignedTask(new AssignedTask().setTask(immediateTaskConfig.setNumCpus(2))));

    TaskConfig immediateTaskConfigTwo =
        defaultTask(false)
            .setJob(
                JOB_KEY.newBuilder().setRole(BAZ_ROLE_IDENTITY.getRole()).setName("immediateTwo"))
            .setJobName("immediateTwo")
            .setOwner(BAZ_ROLE_IDENTITY);
    IScheduledTask task3 =
        IScheduledTask.build(
            new ScheduledTask()
                .setAssignedTask(new AssignedTask().setTask(immediateTaskConfigTwo)));

    TaskConfig immediateTaskConfigThree =
        defaultTask(false)
            .setJob(
                JOB_KEY.newBuilder().setRole(BAZ_ROLE_IDENTITY.getRole()).setName("immediateThree"))
            .setJobName("immediateThree")
            .setOwner(BAZ_ROLE_IDENTITY);
    IScheduledTask task4 =
        IScheduledTask.build(
            new ScheduledTask()
                .setAssignedTask(new AssignedTask().setTask(immediateTaskConfigThree)));

    expect(storageUtil.taskStore.getJobKeys())
        .andReturn(
            FluentIterable.from(ImmutableSet.of(task1, task2, task3, task4))
                .transform(Tasks::getJob)
                .toSet());
    expect(storageUtil.jobStore.fetchJobs()).andReturn(IJobConfiguration.setFromBuilders(crons));

    RoleSummaryResult expectedResult = new RoleSummaryResult();
    expectedResult.addToSummaries(
        new RoleSummary().setRole(ROLE).setCronJobCount(2).setJobCount(1));
    expectedResult.addToSummaries(
        new RoleSummary().setRole(BAZ_ROLE).setCronJobCount(1).setJobCount(2));

    control.replay();

    Response response = assertOkResponse(thrift.getRoleSummary());
    assertEquals(expectedResult, response.getResult().getRoleSummaryResult());
  }