Testing Hazelcast migration with/without backup

package net.peierls.example.hazelcast.migration;

import com.hazelcast.config.*;
import com.hazelcast.core.*;
import com.hazelcast.partition.*;

import java.util.concurrent.CountDownLatch;

import org.junit.*;
import static org.junit.Assert.*;


public class MigrationTest {

    @After public void shutdown() {
        if (stopListeningForMigration != null) {
            stopListeningForMigration.run();
            stopListeningForMigration = null;
        }
        Hazelcast.shutdownAll();
    }


    @Test public void testMigrationNoBackups() {
        doTestMigration(0);
    }

    @Test public void testMigrationOneBackup() {
        doTestMigration(1);
    }


    void doTestMigration(int backupCount) {
        String testId = "test" + backupCount;

        HazelcastInstance h1 = newHazelcastInstance(backupCount);
        HazelcastInstance h2 = newHazelcastInstance(backupCount);
        IMap<Integer, String> m1 = getTestMap(h1);
        IMap<Integer, String> m2 = getTestMap(h2);
        populate(m1, 0, 100);
        populate(m2, 100, 200);

        assertAllValuesPresent(m1);
        assertAllValuesPresent(m2);

        int s1 = m1.localKeySet().size();
        int s2 = m2.localKeySet().size();
        System.out.printf(
            "%n%n%s: m1 has %d local keys, m2 has %d local keys%n%n",
            testId, s1, s2);
        assertEquals(200, s1 + s2);

        listenForMigration(h2.getPartitionService(), testId);
        awaitShutdown(h1);

        System.out.printf(
            "%n%s: m2 has %d keys after shutdown%n%n",
            testId, m2.keySet().size());

        assertAllValuesPresent(m2);
    }


    void populate(IMap<Integer, String> m, int lo, int hi) {
        for (int i = lo; i < hi; ++i) {
            m.put(i, keyToValue(i));
        }
    }

    private void assertAllValuesPresent(IMap<Integer, String> m) {
        for (int i = 0; i < 200; ++i) {
            String s = m.get(i);
            assertNotNull("null value for " + i, s);
            assertEquals(keyToValue(i), s);
        }
    }

    private String keyToValue(int i) {
        return Integer.valueOf(i).toString();
    }

    HazelcastInstance newHazelcastInstance(int backupCount) {
        Config config = new Config();
        config.getMapConfig(TEST_MAP).setBackupCount(backupCount);
        return Hazelcast.newHazelcastInstance(config);
    }

    IMap<Integer, String> getTestMap(HazelcastInstance h) {
        return h.getMap(TEST_MAP);
    }

    void awaitShutdown(HazelcastInstance h) {
        final CountDownLatch latch = new CountDownLatch(1);

        LifecycleService ls = h.getLifecycleService();
        ls.addLifecycleListener(new LifecycleListener() {
            public void stateChanged(LifecycleEvent event) {
                switch (event.getState()) {
                    case SHUTDOWN:
                        latch.countDown();
                }
            }
        });
        ls.shutdown();
        try {
            latch.await();
        } catch (InterruptedException ex) {
            Thread.currentThread().interrupt();
            fail("Unexpected interruption");
        }
    }

    void listenForMigration(final PartitionService ps, final String id) {
        final MigrationListener ml = new MigrationListener() {
            public void migrationStarted(MigrationEvent event) {
                System.out.printf("%s: started %s%n", id, event);
            }
            public void migrationCompleted(MigrationEvent event) {
                System.out.printf("%s: completed %s%n", id, event);
            }
        };
        ps.addMigrationListener(ml);
        stopListeningForMigration = new Runnable() {
            public void run() {
                ps.removeMigrationListener(ml);
            }
        };
    }

    Runnable stopListeningForMigration = null;

    static final String TEST_MAP = "test";
}