@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 upgradeWithFailedStreamUpdate() throws Exception {
        final Stream stream1 = mock(Stream.class);
        final Stream stream2 = mock(Stream.class);
        final IndexSetConfig indexSetConfig = mock(IndexSetConfig.class);

        when(indexSetService.findAll()).thenReturn(Collections.singletonList(indexSetConfig));
        when(indexSetConfig.id()).thenReturn("abc123");
        when(stream1.getId()).thenReturn("stream1");
        when(stream2.getId()).thenReturn("stream2");
        when(streamService.loadAll()).thenReturn(Lists.newArrayList(stream1, stream2));

        // Updating stream1 should fail!
        when(streamService.save(stream1)).thenThrow(ValidationException.class);

        migration.upgrade();

        verify(stream1).setIndexSetId(indexSetConfig.id());
        verify(stream2).setIndexSetId(indexSetConfig.id());
        verify(streamService, times(1)).save(stream1);
        verify(streamService, times(1)).save(stream2);

        // Check that the failed stream1 will be recorded as failed!
        verify(clusterConfigService, times(1)).write(
                V20161122174500_AssignIndexSetsToStreamsMigration.MigrationCompleted.create(
                        indexSetConfig.id(), Sets.newHashSet("stream2"), Sets.newHashSet("stream1")));
    }