Source:
private void testGetNextTrainingDataWithBlockRemovals(final List<Pair<Integer, Map<Long, Integer>>> blocksToRemove) { final int numBlocksToRemove = blocksToRemove.size(); final MoveHandler moveHandler = (MoveHandler) memoryStore; final Map<Long, Integer> trainingDataInStore = new HashMap<>(); final Map<Long, Integer> usedTrainingData = new HashMap<>(); final Map<Long, Integer> removedTrainingData = new HashMap<>(); trainingDataInStore.putAll(kvMapBackingMemoryStore); final int expectedTotalInstances = trainingDataInStore.size(); int miniBatchCount = 0; int removedBlockCount = 0; int numRemovedInstances = 0; int numUsedInstances = 0; Map<Long, Integer> trainingData = trainingDataProvider.getNextTrainingData(); while (!trainingData.isEmpty()) { miniBatchCount++; final Set<Long> keySet = trainingData.keySet(); numUsedInstances += keySet.size(); for (final Long key : keySet) { final Integer expectedValue = trainingDataInStore.remove(key); final Integer actualValue = trainingData.get(key); assertNotNull(expectedValue); assertEquals(expectedValue, actualValue); usedTrainingData.put(key, actualValue); } if (((miniBatchCount % 4) == 0) && (removedBlockCount < numBlocksToRemove)) { final Pair<Integer, Map<Long, Integer>> blockDataPair = blocksToRemove.get(removedBlockCount++); final int blockId = blockDataPair.getFirst(); final Map<Long, Integer> dataSet = blockDataPair.getSecond(); moveHandler.removeBlock(blockId); operationRouter.updateOwnership(blockId, LOCAL_STORE_ID, REMOTE_STORE_ID); for (final Long key : dataSet.keySet()) { final Integer expectedValue = dataSet.get(key); final Integer actualValue = trainingDataInStore.remove(key); if (actualValue != null) { assertEquals(expectedValue, actualValue); removedTrainingData.put(key, actualValue); numRemovedInstances++; } else { assertTrue(usedTrainingData.containsKey(key)); continue; } } } trainingData = trainingDataProvider.getNextTrainingData(); } assertEquals(expectedTotalInstances, numUsedInstances + numRemovedInstances); assertTrue(trainingDataInStore.isEmpty()); }

Target:
private void testGetNextTrainingDataWithBlockRemovals(final List<Pair<Integer, Map<Long, Integer>>> blocksToRemove) { final int numBlocksToRemove = blocksToRemove.size(); final Map<Long, Integer> unusedDataInStore = new HashMap<>(); unusedDataInStore.putAll(memoryStore.getAll()); final int numTotalInstances = unusedDataInStore.size(); int miniBatchIdx = 0; int removedBlockCount = 0; int numRemovedInstances = 0; int numUsedInstances = 0; Map<Long, Integer> trainingData = trainingDataProvider.getNextTrainingData(); while (!trainingData.isEmpty()) { final Set<Long> keySet = trainingData.keySet(); numUsedInstances += keySet.size(); for (final Long key : keySet) { final Integer expectedValue = unusedDataInStore.remove(key); final Integer actualValue = trainingData.get(key); assertNotNull(expectedValue); assertEquals(expectedValue, actualValue); } if (((miniBatchIdx % 4) == 0) && (removedBlockCount < numBlocksToRemove)) { final Pair<Integer, Map<Long, Integer>> blockDataPair = blocksToRemove.get(removedBlockCount); final int blockId = blockDataPair.getFirst(); final Map<Long, Integer> dataSet = blockDataPair.getSecond(); removeBlockFromMemoryStore(blockId); removedBlockCount++; for (final Long key : dataSet.keySet()) { final Integer expectedValue = dataSet.get(key); final Integer actualValue = unusedDataInStore.remove(key); if (actualValue != null) { assertEquals(expectedValue, actualValue); numRemovedInstances++; } } } trainingData = trainingDataProvider.getNextTrainingData(); miniBatchIdx++; } assertEquals(numTotalInstances, numUsedInstances + numRemovedInstances); assertTrue(unusedDataInStore.isEmpty()); }