ImmutableSet<IScheduledTask> tasks = ImmutableSet.of(
IScheduledTask.build(new ScheduledTask().setStatus(ScheduleStatus.PENDING)));
Set<QuotaConfiguration> quotas =
ImmutableSet.of(
new QuotaConfiguration("steve", ResourceAggregates.none().newBuilder()));
IHostAttributes attribute = IHostAttributes.build(new HostAttributes("host",
ImmutableSet.of(new Attribute("attr", ImmutableSet.of("value")))));
StoredJob job = new StoredJob(
"jobManager",
new JobConfiguration().setKey(new JobKey("owner", "env", "name")));
String frameworkId = "framework_id";
ILock lock = ILock.build(new Lock()
.setKey(LockKey.job(JobKeys.from("testRole", "testEnv", "testJob").newBuilder()))
.setToken("lockId")
.setUser("testUser")
.setTimestampMs(12345L));
SchedulerMetadata metadata = new SchedulerMetadata()
.setFrameworkId(frameworkId)
.setVersion(CURRENT_API_VERSION);
storageUtil.expectOperations();
expect(storageUtil.taskStore.fetchTasks(Query.unscoped())).andReturn(tasks);
expect(storageUtil.quotaStore.fetchQuotas())
.andReturn(ImmutableMap.of("steve", ResourceAggregates.none()));
expect(storageUtil.attributeStore.getHostAttributes()).andReturn(ImmutableSet.of(attribute));
expect(storageUtil.jobStore.fetchManagerIds()).andReturn(ImmutableSet.of("jobManager"));
expect(storageUtil.jobStore.fetchJobs("jobManager"))
.andReturn(ImmutableSet.of(IJobConfiguration.build(job.getJobConfiguration())));
expect(storageUtil.schedulerStore.fetchFrameworkId()).andReturn(Optional.of(frameworkId));
expect(storageUtil.lockStore.fetchLocks()).andReturn(ImmutableSet.of(lock));
expectDataWipe();
storageUtil.taskStore.saveTasks(tasks);
storageUtil.quotaStore.saveQuota("steve", ResourceAggregates.none());
storageUtil.attributeStore.saveHostAttributes(attribute);
storageUtil.jobStore.saveAcceptedJob(
job.getJobManagerId(),
IJobConfiguration.build(job.getJobConfiguration()));
storageUtil.schedulerStore.saveFrameworkId(frameworkId);
storageUtil.lockStore.saveLock(lock);
control.replay();
Snapshot expected = new Snapshot()
.setTimestamp(NOW)
.setTasks(IScheduledTask.toBuildersSet(tasks))
.setQuotaConfigurations(quotas)
.setHostAttributes(ImmutableSet.of(attribute.newBuilder()))
.setJobs(ImmutableSet.of(job))
.setSchedulerMetadata(metadata)
.setLocks(ILock.toBuildersSet(ImmutableSet.of(lock)));
assertEquals(expected, snapshotStore.createSnapshot());