@Override
    public int choosePartition(Message<?> msg, TopicMetadata topicMetadata) {
        // If the message has a key, it supersedes the round robin routing policy
        if (msg.hasKey()) {
            return signSafeMod(hash.makeHash(msg.getKey()), topicMetadata.numPartitions());
        }

        if (isBatchingEnabled) { // if batching is enabled, choose partition on `partitionSwitchMs` boundary.
            long currentMs = clock.millis();
            return signSafeMod(currentMs / partitionSwitchMs + startPtnIdx, topicMetadata.numPartitions());
        } else {
            return signSafeMod(PARTITION_INDEX_UPDATER.getAndIncrement(this), topicMetadata.numPartitions());
        }
    }
@Test
    public void testChoosePartitionWithoutKeyWithBatching() {
        Message<?> msg = mock(Message.class);
        Mockito.when(msg.getKey()).thenReturn(null);

        // Fake clock, simulate 1 millisecond passes for each invocation
        Clock clock = new Clock() {
            private long current = 0;

            @Override
            public Clock withZone(ZoneId zone) {
                return null;
            }

            @Override
            public long millis() {
                return current++;
            }

            @Override
            public Instant instant() {
                return Instant.ofEpochMilli(millis());
            }

            @Override
            public ZoneId getZone() {
                return ZoneId.systemDefault();
            }
        };

        RoundRobinPartitionMessageRouterImpl router = new RoundRobinPartitionMessageRouterImpl(
                HashingScheme.JavaStringHash, 0, true, 5, clock);

        // Since the batching time is 5millis, first 5 messages will go on partition 0 and next five would go on
        // partition 1
        for (int i = 0; i < 5; i++) {
            assertEquals(0, router.choosePartition(msg, new TopicMetadataImpl(5)));
        }

        for (int i = 5; i < 10; i++) {
            assertEquals(1, router.choosePartition(msg, new TopicMetadataImpl(5)));
        }
    }