@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); }