@Test public void testTimerAssignmentToKeyGroups() { int totalNoOfTimers = 100; int totalNoOfKeyGroups = 100; int startKeyGroupIdx = 0; int endKeyGroupIdx = totalNoOfKeyGroups - 1; // we have 0 to 99 @SuppressWarnings("unchecked") Set<InternalTimer<Integer, String>>[] expectedNonEmptyTimerSets = new HashSet[totalNoOfKeyGroups]; TestKeyContext keyContext = new TestKeyContext(); HeapInternalTimerService<Integer, String> timerService = new HeapInternalTimerService<>( totalNoOfKeyGroups, new KeyGroupRange(startKeyGroupIdx, endKeyGroupIdx), keyContext, new TestProcessingTimeService()); timerService.startTimerService( IntSerializer.INSTANCE, StringSerializer.INSTANCE, mock(Triggerable.class)); for (int i = 0; i < totalNoOfTimers; i++) { // create the timer to be registered InternalTimer<Integer, String> timer = new InternalTimer<>(10 + i, i, "hello_world_" + i); int keyGroupIdx = KeyGroupRangeAssignment.assignToKeyGroup(timer.getKey(), totalNoOfKeyGroups); // add it in the adequate expected set of timers per keygroup Set<InternalTimer<Integer, String>> timerSet = expectedNonEmptyTimerSets[keyGroupIdx]; if (timerSet == null) { timerSet = new HashSet<>(); expectedNonEmptyTimerSets[keyGroupIdx] = timerSet; } timerSet.add(timer); // register the timer as both processing and event time one keyContext.setCurrentKey(timer.getKey()); timerService.registerEventTimeTimer(timer.getNamespace(), timer.getTimestamp()); timerService.registerProcessingTimeTimer(timer.getNamespace(), timer.getTimestamp()); } Set<InternalTimer<Integer, String>>[] eventTimeTimers = timerService.getEventTimeTimersPerKeyGroup(); Set<InternalTimer<Integer, String>>[] processingTimeTimers = timerService.getProcessingTimeTimersPerKeyGroup(); // finally verify that the actual timers per key group sets are the expected ones. for (int i = 0; i < expectedNonEmptyTimerSets.length; i++) { Set<InternalTimer<Integer, String>> expected = expectedNonEmptyTimerSets[i]; Set<InternalTimer<Integer, String>> actualEvent = eventTimeTimers[i]; Set<InternalTimer<Integer, String>> actualProcessing = processingTimeTimers[i]; if (expected == null) { Assert.assertNull(actualEvent); Assert.assertNull(actualProcessing); } else { Assert.assertArrayEquals(expected.toArray(), actualEvent.toArray()); Assert.assertArrayEquals(expected.toArray(), actualProcessing.toArray()); } } }