| /** |
| * 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.util.Date; |
| import java.util.List; |
| import java.util.Optional; |
| import java.util.Set; |
| import java.util.stream.Collectors; |
| |
| import javax.annotation.Nullable; |
| |
| import com.google.common.base.Function; |
| import com.google.common.collect.FluentIterable; |
| import com.google.common.collect.ImmutableList; |
| import com.google.common.collect.ImmutableSet; |
| import com.google.common.collect.Iterables; |
| import com.google.common.collect.Lists; |
| 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.ConfigGroup; |
| import org.apache.aurora.gen.ConfigSummary; |
| import org.apache.aurora.gen.ConfigSummaryResult; |
| import org.apache.aurora.gen.GetJobUpdateDiffResult; |
| import org.apache.aurora.gen.GetQuotaResult; |
| import org.apache.aurora.gen.GetTierConfigResult; |
| import org.apache.aurora.gen.Identity; |
| import org.apache.aurora.gen.JobConfiguration; |
| import org.apache.aurora.gen.JobKey; |
| import org.apache.aurora.gen.JobStats; |
| import org.apache.aurora.gen.JobSummary; |
| import org.apache.aurora.gen.JobSummaryResult; |
| import org.apache.aurora.gen.JobUpdate; |
| import org.apache.aurora.gen.JobUpdateDetails; |
| import org.apache.aurora.gen.JobUpdateKey; |
| import org.apache.aurora.gen.JobUpdateQuery; |
| import org.apache.aurora.gen.JobUpdateRequest; |
| import org.apache.aurora.gen.JobUpdateSettings; |
| import org.apache.aurora.gen.JobUpdateSummary; |
| import org.apache.aurora.gen.Metadata; |
| import org.apache.aurora.gen.PendingReason; |
| import org.apache.aurora.gen.PopulateJobResult; |
| import org.apache.aurora.gen.Range; |
| import org.apache.aurora.gen.ReadOnlyScheduler; |
| import org.apache.aurora.gen.Response; |
| import org.apache.aurora.gen.ResponseDetail; |
| import org.apache.aurora.gen.Result; |
| import org.apache.aurora.gen.RoleSummary; |
| import org.apache.aurora.gen.RoleSummaryResult; |
| import org.apache.aurora.gen.ScheduleStatus; |
| import org.apache.aurora.gen.ScheduledTask; |
| import org.apache.aurora.gen.TaskConfig; |
| import org.apache.aurora.gen.TaskQuery; |
| import org.apache.aurora.scheduler.TierManager; |
| import org.apache.aurora.scheduler.base.JobKeys; |
| import org.apache.aurora.scheduler.base.Query; |
| import org.apache.aurora.scheduler.base.Query.Builder; |
| import org.apache.aurora.scheduler.base.TaskGroupKey; |
| import org.apache.aurora.scheduler.base.TaskTestUtil; |
| import org.apache.aurora.scheduler.base.Tasks; |
| import org.apache.aurora.scheduler.configuration.SanitizedConfiguration; |
| import org.apache.aurora.scheduler.cron.CronPredictor; |
| import org.apache.aurora.scheduler.cron.CrontabEntry; |
| import org.apache.aurora.scheduler.filter.SchedulingFilter.Veto; |
| import org.apache.aurora.scheduler.metadata.NearestFit; |
| import org.apache.aurora.scheduler.quota.QuotaInfo; |
| import org.apache.aurora.scheduler.quota.QuotaManager; |
| import org.apache.aurora.scheduler.storage.entities.IConfigSummaryResult; |
| import org.apache.aurora.scheduler.storage.entities.IJobConfiguration; |
| import org.apache.aurora.scheduler.storage.entities.IJobKey; |
| import org.apache.aurora.scheduler.storage.entities.IJobUpdateDetails; |
| import org.apache.aurora.scheduler.storage.entities.IJobUpdateQuery; |
| import org.apache.aurora.scheduler.storage.entities.IRange; |
| import org.apache.aurora.scheduler.storage.entities.IResponse; |
| 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.junit.Before; |
| import org.junit.Test; |
| |
| 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.scheduler.base.Numbers.convertRanges; |
| import static org.apache.aurora.scheduler.base.Numbers.toRanges; |
| import static org.apache.aurora.scheduler.resources.ResourceBag.LARGE; |
| import static org.apache.aurora.scheduler.resources.ResourceBag.MEDIUM; |
| import static org.apache.aurora.scheduler.resources.ResourceBag.SMALL; |
| import static org.apache.aurora.scheduler.resources.ResourceBag.XLARGE; |
| import static org.apache.aurora.scheduler.resources.ResourceManager.aggregateFromBag; |
| import static org.apache.aurora.scheduler.thrift.Fixtures.CRON_JOB; |
| import static org.apache.aurora.scheduler.thrift.Fixtures.CRON_SCHEDULE; |
| import static org.apache.aurora.scheduler.thrift.Fixtures.IDENTITY; |
| import static org.apache.aurora.scheduler.thrift.Fixtures.JOB_KEY; |
| import static org.apache.aurora.scheduler.thrift.Fixtures.QUOTA; |
| import static org.apache.aurora.scheduler.thrift.Fixtures.ROLE; |
| 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.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.jobSummaryResponse; |
| import static org.apache.aurora.scheduler.thrift.Fixtures.makeDefaultScheduledTasks; |
| import static org.apache.aurora.scheduler.thrift.Fixtures.makeJob; |
| import static org.apache.aurora.scheduler.thrift.Fixtures.nonProductionTask; |
| import static org.apache.aurora.scheduler.thrift.ReadOnlySchedulerImpl.NO_CRON; |
| import static org.easymock.EasyMock.expect; |
| import static org.junit.Assert.assertEquals; |
| import static org.junit.Assert.assertFalse; |
| import static org.junit.Assert.assertTrue; |
| |
| public class ReadOnlySchedulerImplTest extends EasyMockTest { |
| private static final ImmutableSet<Metadata> METADATA = |
| ImmutableSet.of(new Metadata("k1", "v1"), new Metadata("k2", "v2"), new Metadata("k3", "v3")); |
| |
| private StorageTestUtil storageUtil; |
| private NearestFit nearestFit; |
| private CronPredictor cronPredictor; |
| private QuotaManager quotaManager; |
| private TierManager tierManager; |
| |
| private ReadOnlyScheduler.Iface thrift; |
| |
| @Before |
| public void setUp() { |
| storageUtil = new StorageTestUtil(this); |
| storageUtil.expectOperations(); |
| nearestFit = createMock(NearestFit.class); |
| cronPredictor = createMock(CronPredictor.class); |
| quotaManager = createMock(QuotaManager.class); |
| tierManager = createMock(TierManager.class); |
| |
| thrift = new ReadOnlySchedulerImpl( |
| TaskTestUtil.CONFIGURATION_MANAGER, |
| storageUtil.storage, |
| nearestFit, |
| cronPredictor, |
| quotaManager, |
| tierManager); |
| } |
| |
| @Test |
| public void testGetJobSummary() throws Exception { |
| long nextCronRunMs = 100; |
| TaskConfig ownedCronJobTask = nonProductionTask().setJob(JOB_KEY.newBuilder()); |
| 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().setUser("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")); |
| 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(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 testGetJobSummaryWithoutNextRun() throws Exception { |
| // 31st of February, there is no such day. |
| String cronSchedule = "* * 31 2 *"; |
| |
| TaskConfig task = nonProductionTask() |
| .setJob(JOB_KEY.newBuilder()) |
| .setOwner(IDENTITY); |
| JobConfiguration job = makeJob() |
| .setCronSchedule(cronSchedule) |
| .setTaskConfig(task); |
| expect(cronPredictor.predictNextRun(CrontabEntry.parse(cronSchedule))) |
| .andReturn(Optional.empty()) |
| .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 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().newBuilder())); |
| assertEquals(expected, response.getResult().getGetPendingReasonResult().getReasons()); |
| } |
| |
| @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 testPopulateJobConfigFails() throws Exception { |
| IJobConfiguration job = IJobConfiguration.build(makeJob(null)); |
| control.replay(); |
| |
| assertResponse(INVALID_REQUEST, thrift.populateJobConfig(job.newBuilder())); |
| } |
| |
| @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(aggregateFromBag(QUOTA).newBuilder()) |
| .setProdSharedConsumption(aggregateFromBag(XLARGE).newBuilder()) |
| .setProdDedicatedConsumption(aggregateFromBag(LARGE).newBuilder()) |
| .setNonProdSharedConsumption(aggregateFromBag(MEDIUM).newBuilder()) |
| .setNonProdDedicatedConsumption(aggregateFromBag(SMALL).newBuilder()); |
| |
| Response response = assertOkResponse(thrift.getQuota(ROLE)); |
| assertEquals(expected, response.getResult().getGetQuotaResult()); |
| } |
| |
| @Test |
| public void testGetTasksWithoutConfigs() throws Exception { |
| Builder query = Query.unscoped(); |
| storageUtil.expectTaskFetch(query, ImmutableSet.copyOf(makeDefaultScheduledTasks(10))); |
| |
| control.replay(); |
| |
| TaskConfig taskConfig = defaultTask(true); |
| taskConfig.getExecutorConfig().unsetData(); |
| |
| ImmutableList<ScheduledTask> expected = IScheduledTask.toBuildersList(makeDefaultScheduledTasks( |
| 10, |
| taskConfig)); |
| |
| Response response = assertOkResponse(thrift.getTasksWithoutConfigs(new TaskQuery())); |
| assertEquals(expected, response.getResult().getScheduleStatusResult().getTasks()); |
| } |
| |
| @Test |
| public void testGetPendingReasonFailsSlavesSet() throws Exception { |
| Builder query = Query.unscoped().bySlave("host1"); |
| |
| control.replay(); |
| |
| assertResponse(INVALID_REQUEST, thrift.getPendingReason(query.get().newBuilder())); |
| } |
| |
| @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")); |
| 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(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 testGetJobs() throws Exception { |
| TaskConfig ownedCronJobTask = nonProductionTask(); |
| 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().setUser("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")); |
| 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(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())); |
| } |
| |
| private static void assertJobsEqual(JobConfiguration expected, JobConfiguration actual) { |
| assertEquals(IJobConfiguration.build(expected), IJobConfiguration.build(actual)); |
| } |
| |
| @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); |
| } |
| |
| private TaskQuery setupPaginatedQuery(Iterable<IScheduledTask> tasks, int offset, int limit) { |
| TaskQuery query = new TaskQuery().setOffset(offset).setLimit(limit); |
| Builder builder = Query.arbitrary(query); |
| storageUtil.expectTaskFetch(builder, ImmutableSet.copyOf(tasks)); |
| return query; |
| } |
| |
| private static final Function<ScheduledTask, Integer> TO_INSTANCE_ID = |
| new Function<ScheduledTask, Integer>() { |
| @Nullable |
| @Override |
| public Integer apply(@Nullable ScheduledTask input) { |
| return input.getAssignedTask().getInstanceId(); |
| } |
| }; |
| |
| @Test |
| public void testGetConfigSummary() throws Exception { |
| IJobKey key = JobKeys.from("test", "test", "test"); |
| |
| TaskConfig firstGroupTask = defaultTask(true); |
| TaskConfig secondGroupTask = defaultTask(true) |
| .setResources(ImmutableSet.of(numCpus(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 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 testGetPendingReasonFailsStatusSet() throws Exception { |
| Builder query = Query.unscoped().byStatus(ScheduleStatus.ASSIGNED); |
| |
| control.replay(); |
| |
| assertResponse(INVALID_REQUEST, thrift.getPendingReason(query.get().newBuilder())); |
| } |
| |
| @Test |
| public void testGetJobUpdateSummaries() throws Exception { |
| JobUpdateQuery query = new JobUpdateQuery().setRole(ROLE); |
| Set<JobUpdateDetails> details = createJobUpdateDetails(5); |
| expect(storageUtil.jobUpdateStore.fetchJobUpdates(IJobUpdateQuery.build(query))) |
| .andReturn(IJobUpdateDetails.listFromBuilders(details)); |
| |
| control.replay(); |
| |
| Response response = assertOkResponse(thrift.getJobUpdateSummaries(query)); |
| assertEquals( |
| details.stream().map(u -> u.getUpdate().getSummary()).collect(Collectors.toList()), |
| response.getResult().getGetJobUpdateSummariesResult().getUpdateSummaries()); |
| } |
| |
| @Test |
| public void testGetJobUpdateDetails() throws Exception { |
| JobUpdateDetails details = createJobUpdateDetails(); |
| expect(storageUtil.jobUpdateStore.fetchJobUpdate(UPDATE_KEY)) |
| .andReturn(Optional.of(IJobUpdateDetails.build(details))); |
| |
| control.replay(); |
| |
| Response response = assertOkResponse(thrift.getJobUpdateDetails(UPDATE_KEY.newBuilder(), null)); |
| assertEquals( |
| IJobUpdateDetails.build(details), |
| IJobUpdateDetails.build(response.getResult().getGetJobUpdateDetailsResult().getDetails())); |
| // Specifying the UPDATE_KEY alone is deprecated, so there should be a message. |
| assertEquals(1, response.getDetailsSize()); |
| } |
| |
| @Test |
| public void testGetJobUpdateDetailsQuery() throws Exception { |
| JobUpdateQuery query = new JobUpdateQuery().setRole(ROLE); |
| List<IJobUpdateDetails> details = IJobUpdateDetails.listFromBuilders(createJobUpdateDetails(5)); |
| expect(storageUtil.jobUpdateStore.fetchJobUpdates(IJobUpdateQuery.build(query))) |
| .andReturn(details); |
| |
| control.replay(); |
| |
| Response response = assertOkResponse(thrift.getJobUpdateDetails(null, query)); |
| assertEquals( |
| IJobUpdateDetails.toBuildersList(details), |
| response.getResult().getGetJobUpdateDetailsResult().getDetailsList()); |
| } |
| |
| private static List<JobUpdateSummary> createJobUpdateSummaries(int count) { |
| ImmutableList.Builder<JobUpdateSummary> builder = ImmutableList.builder(); |
| for (int i = 0; i < count; i++) { |
| builder.add(new JobUpdateSummary() |
| .setKey(new JobUpdateKey(JOB_KEY.newBuilder(), "id" + 1)) |
| .setUser(USER) |
| .setMetadata(METADATA)); |
| } |
| return builder.build(); |
| } |
| |
| private static Set<JobUpdateDetails> createJobUpdateDetails(int count) { |
| List<JobUpdateSummary> summaries = createJobUpdateSummaries(count); |
| return summaries.stream() |
| .map(jobUpdateSummary -> |
| new JobUpdateDetails().setUpdate(new JobUpdate().setSummary(jobUpdateSummary))) |
| .collect(Collectors.toSet()); |
| } |
| |
| private static JobUpdateDetails createJobUpdateDetails() { |
| return createJobUpdateDetails(1).stream().findFirst().get(); |
| } |
| |
| @Test |
| public void testGetRoleSummary() throws Exception { |
| String bazRole = "baz_role"; |
| Identity bazRoleIdentity = new Identity().setUser(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(bazRole)) |
| .setTaskConfig(nonProductionTask()) |
| .setOwner(bazRoleIdentity); |
| |
| Set<JobConfiguration> crons = ImmutableSet.of(cronJobOne, cronJobTwo, cronJobThree); |
| |
| TaskConfig immediateTaskConfig = defaultTask(false) |
| .setJob(JOB_KEY.newBuilder().setName("immediate")); |
| IScheduledTask task1 = IScheduledTask.build(new ScheduledTask() |
| .setAssignedTask(new AssignedTask().setTask(immediateTaskConfig))); |
| IScheduledTask task2 = IScheduledTask.build(new ScheduledTask() |
| .setAssignedTask(new AssignedTask().setTask(immediateTaskConfig |
| .setResources(ImmutableSet.of( |
| numCpus(2), |
| ramMb(1024), |
| diskMb(1024)))))); |
| |
| TaskConfig immediateTaskConfigTwo = defaultTask(false) |
| .setJob(JOB_KEY.newBuilder().setRole(bazRole).setName("immediateTwo")) |
| .setOwner(bazRoleIdentity); |
| IScheduledTask task3 = IScheduledTask.build(new ScheduledTask() |
| .setAssignedTask(new AssignedTask().setTask(immediateTaskConfigTwo))); |
| |
| TaskConfig immediateTaskConfigThree = defaultTask(false) |
| .setJob(JOB_KEY.newBuilder().setRole(bazRole).setName("immediateThree")) |
| .setOwner(bazRoleIdentity); |
| 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(bazRole).setCronJobCount(1).setJobCount(2)); |
| |
| control.replay(); |
| |
| Response response = assertOkResponse(thrift.getRoleSummary()); |
| assertEquals(expectedResult, response.getResult().getRoleSummaryResult()); |
| } |
| |
| @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 testGetJobUpdateDetailsInvalidId() throws Exception { |
| expect(storageUtil.jobUpdateStore.fetchJobUpdate(UPDATE_KEY)) |
| .andReturn(Optional.empty()); |
| |
| control.replay(); |
| |
| assertResponse(INVALID_REQUEST, thrift.getJobUpdateDetails(UPDATE_KEY.newBuilder(), null)); |
| } |
| |
| @Test |
| public void testGetJobUpdateDiffWithUpdateAdd() throws Exception { |
| TaskConfig task1 = defaultTask(false).setMaxTaskFailures(1); |
| TaskConfig task2 = defaultTask(false).setMaxTaskFailures(2); |
| TaskConfig task3 = defaultTask(false).setMaxTaskFailures(3); |
| TaskConfig task4 = defaultTask(false).setMaxTaskFailures(4); |
| TaskConfig task5 = defaultTask(false).setMaxTaskFailures(5); |
| |
| ImmutableSet.Builder<IScheduledTask> tasks = ImmutableSet.builder(); |
| makeTasks(0, 10, task1, tasks); |
| makeTasks(10, 20, task2, tasks); |
| makeTasks(20, 30, task3, tasks); |
| makeTasks(30, 40, task4, tasks); |
| makeTasks(40, 50, task5, tasks); |
| |
| expect(storageUtil.jobStore.fetchJob(JOB_KEY)).andReturn(Optional.empty()); |
| storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active(), tasks.build()); |
| |
| control.replay(); |
| |
| TaskConfig newTask = defaultTask(false).setMaxTaskFailures(6); |
| JobUpdateRequest request = new JobUpdateRequest() |
| .setTaskConfig(newTask) |
| .setInstanceCount(60) |
| .setSettings(new JobUpdateSettings() |
| .setUpdateOnlyTheseInstances(ImmutableSet.of(new Range(10, 59)))); |
| |
| GetJobUpdateDiffResult expected = new GetJobUpdateDiffResult() |
| .setAdd(ImmutableSet.of(group(newTask, new Range(50, 59)))) |
| .setUpdate(ImmutableSet.of( |
| group(task2, new Range(10, 19)), |
| group(task3, new Range(20, 29)), |
| group(task4, new Range(30, 39)), |
| group(task5, new Range(40, 49)))) |
| .setUnchanged(ImmutableSet.of(group(task1, new Range(0, 9)))) |
| .setRemove(ImmutableSet.of()); |
| |
| Response response = assertOkResponse(thrift.getJobUpdateDiff(request)); |
| assertEquals(expected, response.getResult().getGetJobUpdateDiffResult()); |
| } |
| |
| @Test |
| public void testGetJobUpdateDiffWithUpdateRemove() throws Exception { |
| TaskConfig task1 = defaultTask(false).setMaxTaskFailures(1); |
| TaskConfig task2 = defaultTask(false).setMaxTaskFailures(2); |
| TaskConfig task3 = defaultTask(false).setMaxTaskFailures(3); |
| |
| 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.empty()); |
| storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active(), tasks.build()); |
| |
| control.replay(); |
| |
| JobUpdateRequest request = new JobUpdateRequest() |
| .setTaskConfig(defaultTask(false).setMaxTaskFailures(6)) |
| .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 testGetJobUpdateDiffWithUnchanged() throws Exception { |
| expect(storageUtil.jobStore.fetchJob(JOB_KEY)).andReturn(Optional.empty()); |
| 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 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 testGetJobUpdateDiffInvalidConfig() throws Exception { |
| control.replay(); |
| TaskConfig task = defaultTask(false) |
| .setResources(ImmutableSet.of( |
| numCpus(-1), |
| ramMb(1024), |
| diskMb(1024))); |
| |
| JobUpdateRequest request = |
| new JobUpdateRequest().setTaskConfig(task); |
| assertResponse(INVALID_REQUEST, thrift.getJobUpdateDiff(request)); |
| } |
| |
| private static void makeTasks( |
| int start, |
| int end, |
| TaskConfig config, |
| ImmutableSet.Builder<IScheduledTask> builder) { |
| |
| for (int i = start; i < end; i++) { |
| builder.add(IScheduledTask.build(new ScheduledTask() |
| .setAssignedTask(new AssignedTask().setTask(config).setInstanceId(i)))); |
| } |
| } |
| |
| private static ConfigGroup group(TaskConfig task, Range range) { |
| return new ConfigGroup() |
| .setConfig(task) |
| .setInstances(ImmutableSet.of(range)); |
| } |
| |
| @Test |
| public void testGetTierConfig() throws Exception { |
| expect(tierManager.getDefaultTierName()).andReturn("preemptible"); |
| expect(tierManager.getTiers()).andReturn(TaskTestUtil.tierInfos()); |
| control.replay(); |
| |
| GetTierConfigResult expected = new GetTierConfigResult() |
| .setDefaultTierName("preemptible") |
| .setTiers(TaskTestUtil.tierConfigs()); |
| |
| Response response = assertOkResponse(thrift.getTierConfigs()); |
| assertEquals(expected, response.getResult().getGetTierConfigResult()); |
| } |
| } |
