@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 testChoosePartitionWithKey() {
        String key1 = "key1";
        String key2 = "key2";
        Message<?> msg1 = mock(Message.class);
        when(msg1.hasKey()).thenReturn(true);
        when(msg1.getKey()).thenReturn(key1);
        Message<?> msg2 = mock(Message.class);
        when(msg2.hasKey()).thenReturn(true);
        when(msg2.getKey()).thenReturn(key2);

        RoundRobinPartitionMessageRouterImpl router = new RoundRobinPartitionMessageRouterImpl(
                HashingScheme.JavaStringHash, 0, false, 0);
        TopicMetadataImpl metadata = new TopicMetadataImpl(100);

        assertEquals(key1.hashCode() % 100, router.choosePartition(msg1, metadata));
        assertEquals(key2.hashCode() % 100, router.choosePartition(msg2, metadata));
    }