@Override
    public void upgrade() {
        // Only run this migration once.
        if (clusterConfigService.get(MigrationCompleted.class) != null) {
            LOG.debug("Migration already completed.");
            return;
        }

        final IndexSetConfig indexSetConfig = findDefaultIndexSet();
        final ImmutableSet.Builder<String> completedStreamIds = ImmutableSet.builder();
        final ImmutableSet.Builder<String> failedStreamIds = ImmutableSet.builder();

        // Assign the "default index set" to all existing streams. Until now, there was no way to manually create
        // index sets, so the only one that exists is the "default" one created by an earlier migration.
        for (Stream stream : streamService.loadAll()) {
            if (isNullOrEmpty(stream.getIndexSetId())) {
                LOG.info("Assigning index set <{}> ({}) to stream <{}> ({})", indexSetConfig.id(),
                        indexSetConfig.title(), stream.getId(), stream.getTitle());
                stream.setIndexSetId(indexSetConfig.id());
                try {
                    streamService.save(stream);
                    completedStreamIds.add(stream.getId());
                } catch (ValidationException e) {
                    LOG.error("Unable to save stream <{}>", stream.getId(), e);
                    failedStreamIds.add(stream.getId());
                }
            }
        }

        // Mark this migration as done.
        clusterConfigService.write(MigrationCompleted.create(indexSetConfig.id(), completedStreamIds.build(), failedStreamIds.build()));
    }
@Test
    public void upgradeWhenAlreadyCompleted() throws Exception {
        final IndexSetConfig indexSetConfig = mock(IndexSetConfig.class);

        when(indexSetService.findAll()).thenReturn(Collections.singletonList(indexSetConfig));
        when(indexSetConfig.id()).thenReturn("abc123");
        when(clusterConfigService.get(V20161122174500_AssignIndexSetsToStreamsMigration.MigrationCompleted.class))
                .thenReturn(V20161122174500_AssignIndexSetsToStreamsMigration.MigrationCompleted.create("1", Collections.emptySet(), Collections.emptySet()));

        migration.upgrade();

        verify(streamService, never()).save(any(Stream.class));
        verify(clusterConfigService, never()).write(any(V20161122174500_AssignIndexSetsToStreamsMigration.MigrationCompleted.class));
    }