@ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
    @EnumSource(RackConfig.class)
    public void testStickiness(RackConfig rackConfig) {
        initializeRacks(rackConfig);
        Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
        partitionsPerTopic.put(topic1, partitionInfos(topic1, 3));

        subscriptions.put(consumer1, subscription(topics(topic1), 0));
        subscriptions.put(consumer2, subscription(topics(topic1), 1));
        subscriptions.put(consumer3, subscription(topics(topic1), 2));
        subscriptions.put(consumer4, subscription(topics(topic1), 3));

        Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
        assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
        Map<String, TopicPartition> partitionsAssigned = new HashMap<>();

        Set<Map.Entry<String, List<TopicPartition>>> assignments = assignment.entrySet();
        for (Map.Entry<String, List<TopicPartition>> entry: assignments) {
            String consumer = entry.getKey();
            List<TopicPartition> topicPartitions = entry.getValue();
            int size = topicPartitions.size();
            assertTrue(size <= 1, "Consumer " + consumer + " is assigned more topic partitions than expected.");
            if (size == 1)
                partitionsAssigned.put(consumer, topicPartitions.get(0));
        }

        // removing the potential group leader
        subscriptions.remove(consumer1);
        subscriptions.put(consumer2,
            buildSubscriptionV2Above(topics(topic1), assignment.get(consumer2), generationId, 1));
        subscriptions.put(consumer3,
            buildSubscriptionV2Above(topics(topic1), assignment.get(consumer3), generationId, 2));
        subscriptions.put(consumer4,
            buildSubscriptionV2Above(topics(topic1), assignment.get(consumer4), generationId, 3));


        assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
        assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);

        assignments = assignment.entrySet();
        for (Map.Entry<String, List<TopicPartition>> entry: assignments) {
            String consumer = entry.getKey();
            List<TopicPartition> topicPartitions = entry.getValue();
            assertEquals(1, topicPartitions.size(), "Consumer " + consumer + " is assigned more topic partitions than expected.");
            assertTrue((!partitionsAssigned.containsKey(consumer)) || (assignment.get(consumer).contains(partitionsAssigned.get(consumer))),
                "Stickiness was not honored for consumer " + consumer);
        }
    }
This unit test performs sticky assignment for a scenario that round robin assignor handles poorly.
Topics (partitions per topic): topic1 (2), topic2 (1), topic3 (2), topic4 (1), topic5 (2)
Subscriptions:
 - consumer1: topic1, topic2, topic3, topic4, topic5
 - consumer2: topic1, topic3, topic5
 - consumer3: topic1, topic3, topic5
 - consumer4: topic1, topic2, topic3, topic4, topic5
Round Robin Assignment Result:
 - consumer1: topic1-0, topic3-0, topic5-0
 - consumer2: topic1-1, topic3-1, topic5-1
 - consumer3:
 - consumer4: topic2-0, topic4-0
Sticky Assignment Result:
 - consumer1: topic2-0, topic3-0
 - consumer2: topic1-0, topic3-1
 - consumer3: topic1-1, topic5-0
 - consumer4: topic4-0, topic5-1
clients/src/test/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignorTest.java