@Test
public void testCreateAndRestoreNewSnapshot() {
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(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());
snapshotStore.applySnapshot(expected);
}
