Source:
<code>public void testMigration() { try { final CountDownLatch waitUpdateLatch = new CountDownLatch(NUM_REQUESTS); final ExecutorService startExecutor = Executors.newFixedThreadPool(NUM_THREAD); for (int i = 0; i < NUM_REQUESTS; i++) { final int index = i; startExecutor.execute(new Runnable() { @Override public void run() { final Set<LongRange> rangesToMove = new HashSet<>(); rangesToMove.add(new LongRange(2 * index, 2 * index + 1)); if (index < 5) { migrationManager.startMigration(index + "", EVAL0, EVAL1, DATA_TYPE, rangesToMove, null); } else { migrationManager.startMigration(index + "", EVAL1, EVAL0, DATA_TYPE, rangesToMove, null); } migrationManager.waitUpdate(index + ""); waitUpdateLatch.countDown(); } }); } waitUpdateLatch.await(); final CountDownLatch updateLatch = new CountDownLatch(1); for (int i = 0; i < NUM_REQUESTS; i++) { final ExecutorService updateExecutor = Executors.newFixedThreadPool(NUM_THREAD); final int index = i; updateExecutor.execute(new Runnable() { @Override public void run() { try { updateLatch.await(); <START> Thread.sleep(10); <END> migrationManager.movePartition(index + "", null); migrationManager.updateSender(index + "", null); migrationManager.finishMigration(index + ""); } catch (final InterruptedException e) { throw new RuntimeException(e); } } }); } updateLatch.countDown(); migrationManager.applyUpdates(null); assertEquals(1, partitionManager.getRangeSet(EVAL0, DATA_TYPE).size()); for (final LongRange range0 : partitionManager.getRangeSet(EVAL0, DATA_TYPE)) { assertEquals(range0, INITIAL_RANGE1); } assertEquals(1, partitionManager.getRangeSet(EVAL0, DATA_TYPE).size()); for (final LongRange range0 : partitionManager.getRangeSet(EVAL0, DATA_TYPE)) { assertEquals(range0, INITIAL_RANGE1); } } catch (final InterruptedException e) { throw new RuntimeException(e); } }</code><technical_language>risk of making this test more elaborate, avoid sleep? idea: introduce another latch applyUpdatesLatch start a new thread below. new thread will: call applyUpdates() call applyUpdatesLatch.countDown(). instead of sleep here, applyUpdatesLatch.await(). this work</technical_language>

Target:
public void testMigration() { try { partitionManager.register(EVAL0, DATA_TYPE, INITIAL_RANGE0); partitionManager.register(EVAL1, DATA_TYPE, INITIAL_RANGE1); final CountDownLatch waitUpdateLatch = new CountDownLatch(NUM_REQUESTS); final ExecutorService startExecutor = Executors.newFixedThreadPool(NUM_THREAD); for (int i = 0; i < NUM_REQUESTS; i++) { final int index = i; startExecutor.execute(new Runnable() { @Override public void run() { final Set<LongRange> rangesToMove = new HashSet<>(); rangesToMove.add(new LongRange(2 * index, 2 * index + 1)); if (index < 5) { migrationManager.startMigration(Integer.toString(index), EVAL0, EVAL1, DATA_TYPE, rangesToMove, null); } else { migrationManager.startMigration(Integer.toString(index), EVAL1, EVAL0, DATA_TYPE, rangesToMove, null); } migrationManager.setMovedRanges(Integer.toString(index), rangesToMove); migrationManager.waitUpdate(Integer.toString(index)); waitUpdateLatch.countDown(); } }); } waitUpdateLatch.await(); final CountDownLatch updateLatch = new CountDownLatch(1); for (int i = 0; i < NUM_REQUESTS; i++) { final ExecutorService updateExecutor = Executors.newFixedThreadPool(NUM_THREAD); final int index = i; updateExecutor.execute(new Runnable() { @Override public void run() { try { updateLatch.await(); Thread.sleep(10); migrationManager.movePartition(Integer.toString(index), null); migrationManager.updateSender(Integer.toString(index), null); migrationManager.finishMigration(Integer.toString(index)); } catch (final InterruptedException e) { throw new RuntimeException(e); } } }); } updateLatch.countDown(); migrationManager.applyUpdates(null); assertEquals(1, partitionManager.getRangeSet(EVAL0, DATA_TYPE).size()); for (final LongRange range0 : partitionManager.getRangeSet(EVAL0, DATA_TYPE)) { assertEquals(range0, INITIAL_RANGE1); } assertEquals(1, partitionManager.getRangeSet(EVAL0, DATA_TYPE).size()); for (final LongRange range0 : partitionManager.getRangeSet(EVAL0, DATA_TYPE)) { assertEquals(range0, INITIAL_RANGE1); } } catch (final InterruptedException e) { throw new RuntimeException(e); } }