content: /*
   * Licensed to the Apache Software Foundation (ASF) under one or more
   * contributor license agreements. See the NOTICE file distributed with
   * this work for additional information regarding copyright ownership.
   * The ASF licenses this file to You under the Apache License, Version 2.0
   * (the "License"); you may not use this file except in compliance with
   * the License. You may obtain a copy of the License at
   *
   *    http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package org.apache.kafka.clients.consumer.internals;
  
  import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.Subscription;
  import org.apache.kafka.clients.consumer.CooperativeStickyAssignor;
  import org.apache.kafka.clients.consumer.StickyAssignor;
  import org.apache.kafka.clients.consumer.internals.AbstractPartitionAssignorTest.RackConfig;
  import org.apache.kafka.common.PartitionInfo;
  import org.apache.kafka.common.TopicPartition;
  import org.apache.kafka.common.test.api.Flaky;
  import org.apache.kafka.common.utils.CollectionUtils;
  import org.apache.kafka.common.utils.Utils;
  
  import org.junit.jupiter.api.BeforeEach;
  import org.junit.jupiter.api.Test;
  import org.junit.jupiter.api.Timeout;
  import org.junit.jupiter.params.ParameterizedTest;
  import org.junit.jupiter.params.provider.EnumSource;
  import org.junit.jupiter.params.provider.ValueSource;
  
  import java.nio.ByteBuffer;
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Map;
  import java.util.Optional;
  import java.util.Random;
  import java.util.Set;
  import java.util.stream.Collectors;
  
  import static java.util.Arrays.asList;
  import static java.util.Collections.emptyList;
  import static org.apache.kafka.clients.consumer.internals.AbstractPartitionAssignorTest.TEST_NAME_WITH_CONSUMER_RACK;
  import static org.apache.kafka.clients.consumer.internals.AbstractPartitionAssignorTest.TEST_NAME_WITH_RACK_CONFIG;
  import static org.apache.kafka.clients.consumer.internals.AbstractPartitionAssignorTest.nullRacks;
  import static org.apache.kafka.clients.consumer.internals.AbstractPartitionAssignorTest.racks;
  import static org.apache.kafka.clients.consumer.internals.AbstractPartitionAssignorTest.verifyRackAssignment;
  import static org.apache.kafka.clients.consumer.internals.AbstractStickyAssignor.DEFAULT_GENERATION;
  import static org.apache.kafka.common.utils.Utils.mkEntry;
  import static org.apache.kafka.common.utils.Utils.mkMap;
  import static org.junit.jupiter.api.Assertions.assertEquals;
  import static org.junit.jupiter.api.Assertions.assertFalse;
  import static org.junit.jupiter.api.Assertions.assertNull;
  import static org.junit.jupiter.api.Assertions.assertTrue;
  
  public abstract class AbstractStickyAssignorTest {
      protected AbstractStickyAssignor assignor;
      protected String consumerId = "consumer";
      protected String consumer1 = "consumer1";
      protected String consumer2 = "consumer2";
      protected String consumer3 = "consumer3";
      protected String consumer4 = "consumer4";
      protected Map<String, Subscription> subscriptions;
      protected String topic = "topic";
      protected String topic1 = "topic1";
      protected String topic2 = "topic2";
      protected String topic3 = "topic3";
      protected TopicPartition tp0 = tp(topic, 0);
      protected TopicPartition tp1 = tp(topic, 1);
      protected TopicPartition tp2 = tp(topic, 2);
      protected String groupId = "group";
      protected int generationId = 1;
  
      protected int numBrokerRacks;
      protected boolean hasConsumerRack;
      protected int replicationFactor = 2;
      private int nextPartitionIndex;
  
      protected abstract AbstractStickyAssignor createAssignor();
  
      // simulate ConsumerProtocolSubscription V0 protocol
      protected abstract Subscription buildSubscriptionV0(List<String> topics, List<TopicPartition> partitions, int generationId, int consumerIndex);
  
      // simulate ConsumerProtocolSubscription V1 protocol
      protected abstract Subscription buildSubscriptionV1(List<String> topics, List<TopicPartition> partitions, int generationId, int consumerIndex);
  
      // simulate ConsumerProtocolSubscription V2 or above protocol
      protected abstract Subscription buildSubscriptionV2Above(List<String> topics, List<TopicPartition> partitions, int generation, int consumerIndex);
  
      protected abstract ByteBuffer generateUserData(List<String> topics, List<TopicPartition> partitions, int generation);
  
      @BeforeEach
      public void setUp() {
          assignor = createAssignor();
  
          if (subscriptions != null) {
              subscriptions.clear();
          } else {
              subscriptions = new HashMap<>();
          }
      }
  
      @Test
      public void testMemberData() {
          List<String> topics = topics(topic);
          List<TopicPartition> ownedPartitions = partitions(tp(topic1, 0), tp(topic2, 1));
          List<Subscription> subscriptions = new ArrayList<>();
          // add subscription in all ConsumerProtocolSubscription versions
          subscriptions.add(buildSubscriptionV0(topics, ownedPartitions, generationId, 0));
          subscriptions.add(buildSubscriptionV1(topics, ownedPartitions, generationId, 1));
          subscriptions.add(buildSubscriptionV2Above(topics, ownedPartitions, generationId, 2));
          for (Subscription subscription : subscriptions) {
              if (subscription != null) {
                  AbstractStickyAssignor.MemberData memberData = assignor.memberData(subscription);
                  assertEquals(ownedPartitions, memberData.partitions, "subscription: " + subscription + " doesn't have expected owned partition");
                  assertEquals(generationId, memberData.generation.orElse(-1), "subscription: " + subscription + " doesn't have expected generation id");
              }
          }
      }
  
      @ParameterizedTest(name = TEST_NAME_WITH_CONSUMER_RACK)
      @ValueSource(booleans = {false, true})
      public void testOneConsumerNoTopic(boolean hasConsumerRack) {
          initializeRacks(hasConsumerRack ? RackConfig.BROKER_AND_CONSUMER_RACK : RackConfig.NO_CONSUMER_RACK);
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          subscriptions = Collections.singletonMap(consumerId, subscription(Collections.emptyList(), 0));
  
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          assertEquals(Collections.singleton(consumerId), assignment.keySet());
          assertTrue(assignment.get(consumerId).isEmpty());
          assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
  
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          assertTrue(isFullyBalanced(assignment));
      }
  
      @ParameterizedTest(name = TEST_NAME_WITH_CONSUMER_RACK)
      @ValueSource(booleans = {false, true})
      public void testOneConsumerNonexistentTopic(boolean hasConsumerRack) {
          initializeRacks(hasConsumerRack ? RackConfig.BROKER_AND_CONSUMER_RACK : RackConfig.NO_CONSUMER_RACK);
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          partitionsPerTopic.put(topic, Collections.emptyList());
          subscriptions = Collections.singletonMap(consumerId, subscription(topics(topic), 0));
  
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
  
          assertEquals(Collections.singleton(consumerId), assignment.keySet());
          assertTrue(assignment.get(consumerId).isEmpty());
          assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
  
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          assertTrue(isFullyBalanced(assignment));
      }
  
      @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
      @EnumSource(RackConfig.class)
      public void testOneConsumerOneTopic(RackConfig rackConfig) {
          initializeRacks(rackConfig);
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          partitionsPerTopic.put(topic, partitionInfos(topic, 3));
          subscriptions = Collections.singletonMap(consumerId, subscription(topics(topic), 0));
  
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          assertEquals(partitions(tp(topic, 0), tp(topic, 1), tp(topic, 2)), assignment.get(consumerId));
          assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
  
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          assertTrue(isFullyBalanced(assignment));
      }
  
      @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
      @EnumSource(RackConfig.class)
      public void testOnlyAssignsPartitionsFromSubscribedTopics(RackConfig rackConfig) {
          String otherTopic = "other";
  
          initializeRacks(rackConfig);
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          partitionsPerTopic.put(topic, partitionInfos(topic, 2));
          subscriptions = mkMap(
                  mkEntry(consumerId, buildSubscriptionV2Above(
                          topics(topic),
                          Arrays.asList(tp(topic, 0), tp(topic, 1), tp(otherTopic, 0), tp(otherTopic, 1)),
                          generationId, 0)
                  )
          );
  
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          assertEquals(partitions(tp(topic, 0), tp(topic, 1)), assignment.get(consumerId));
          assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
  
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          assertTrue(isFullyBalanced(assignment));
      }
  
      @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
      @EnumSource(RackConfig.class)
      public void testOneConsumerMultipleTopics(RackConfig rackConfig) {
          initializeRacks(rackConfig);
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          partitionsPerTopic.put(topic1, partitionInfos(topic1, 1));
          partitionsPerTopic.put(topic2, partitionInfos(topic2, 2));
          subscriptions = Collections.singletonMap(consumerId, subscription(topics(topic1, topic2), 0));
  
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          assertEquals(partitions(tp(topic1, 0), tp(topic2, 0), tp(topic2, 1)), assignment.get(consumerId));
          assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
  
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          assertTrue(isFullyBalanced(assignment));
      }
  
      @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
      @EnumSource(RackConfig.class)
      public void testTwoConsumersOneTopicOnePartition(RackConfig rackConfig) {
          initializeRacks(rackConfig);
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          partitionsPerTopic.put(topic, partitionInfos(topic, 1));
  
          subscriptions.put(consumer1, subscription(topics(topic), 0));
          subscriptions.put(consumer2, subscription(topics(topic), 1));
  
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
  
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          assertTrue(isFullyBalanced(assignment));
      }
  
      @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
      @EnumSource(RackConfig.class)
      public void testTwoConsumersOneTopicTwoPartitions(RackConfig rackConfig) {
          initializeRacks(rackConfig);
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          partitionsPerTopic.put(topic, partitionInfos(topic, 2));
  
          subscriptions.put(consumer1, subscription(topics(topic), 0));
          subscriptions.put(consumer2, subscription(topics(topic), 1));
  
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          assertEquals(partitions(tp(topic, 0)), assignment.get(consumer1));
          assertEquals(partitions(tp(topic, 1)), assignment.get(consumer2));
          assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
  
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          assertTrue(isFullyBalanced(assignment));
      }
  
      @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
      @EnumSource(RackConfig.class)
      public void testMultipleConsumersMixedTopicSubscriptions(RackConfig rackConfig) {
          initializeRacks(rackConfig);
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          partitionsPerTopic.put(topic1, partitionInfos(topic1, 3));
          partitionsPerTopic.put(topic2, partitionInfos(topic2, 2));
  
          subscriptions.put(consumer1, subscription(topics(topic1), 0));
          subscriptions.put(consumer2, subscription(topics(topic1, topic2), 1));
          subscriptions.put(consumer3, subscription(topics(topic1), 2));
  
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          assertEquals(partitions(tp(topic1, 0), tp(topic1, 2)), assignment.get(consumer1));
          assertEquals(partitions(tp(topic2, 0), tp(topic2, 1)), assignment.get(consumer2));
          assertEquals(partitions(tp(topic1, 1)), assignment.get(consumer3));
          assertNull(assignor.partitionsTransferringOwnership);
  
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          assertTrue(isFullyBalanced(assignment));
      }
  
      @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
      @EnumSource(RackConfig.class)
      public void testTwoConsumersTwoTopicsSixPartitions(RackConfig rackConfig) {
          initializeRacks(rackConfig);
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          partitionsPerTopic.put(topic1, partitionInfos(topic1, 3));
          partitionsPerTopic.put(topic2, partitionInfos(topic2, 3));
  
          subscriptions.put(consumer1, subscription(topics(topic1, topic2), 0));
          subscriptions.put(consumer2, subscription(topics(topic1, topic2), 1));
  
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          assertEquals(partitions(tp(topic1, 0), tp(topic1, 2), tp(topic2, 1)), assignment.get(consumer1));
          assertEquals(partitions(tp(topic1, 1), tp(topic2, 0), tp(topic2, 2)), assignment.get(consumer2));
          assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
  
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          assertTrue(isFullyBalanced(assignment));
      }
  
      /**
       * This unit test is testing consumer owned minQuota partitions, and expected to have maxQuota partitions situation
       */
      @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
      @EnumSource(RackConfig.class)
      public void testConsumerOwningMinQuotaExpectedMaxQuota(RackConfig rackConfig) {
          initializeRacks(rackConfig);
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          partitionsPerTopic.put(topic1, partitionInfos(topic1, 2));
          partitionsPerTopic.put(topic2, partitionInfos(topic2, 3));
  
          List<String> subscribedTopics = topics(topic1, topic2);
  
          subscriptions.put(consumer1,
              buildSubscriptionV2Above(subscribedTopics, partitions(tp(topic1, 0), tp(topic2, 1)), generationId, 0));
          subscriptions.put(consumer2,
              buildSubscriptionV2Above(subscribedTopics, partitions(tp(topic1, 1), tp(topic2, 2)), generationId, 1));
  
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          assertEquals(partitions(tp(topic1, 0), tp(topic2, 1), tp(topic2, 0)), assignment.get(consumer1));
          assertEquals(partitions(tp(topic1, 1), tp(topic2, 2)), assignment.get(consumer2));
          assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
  
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          assertTrue(isFullyBalanced(assignment));
      }
  
      /**
       * This unit test is testing consumers owned maxQuota partitions are more than numExpectedMaxCapacityMembers situation
       */
      @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
      @EnumSource(RackConfig.class)
      public void testMaxQuotaConsumerMoreThanNumExpectedMaxCapacityMembers(RackConfig rackConfig) {
          initializeRacks(rackConfig);
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          partitionsPerTopic.put(topic1, partitionInfos(topic1, 2));
          partitionsPerTopic.put(topic2, partitionInfos(topic2, 2));
  
          List<String> subscribedTopics = topics(topic1, topic2);
  
          subscriptions.put(consumer1,
              buildSubscriptionV2Above(subscribedTopics, partitions(tp(topic1, 0), tp(topic2, 0)), generationId, 0));
          subscriptions.put(consumer2,
              buildSubscriptionV2Above(subscribedTopics, partitions(tp(topic1, 1), tp(topic2, 1)), generationId, 1));
          subscriptions.put(consumer3, buildSubscriptionV2Above(subscribedTopics, Collections.emptyList(), generationId, 2));
  
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          assertEquals(Collections.singletonMap(tp(topic2, 0), consumer3), assignor.partitionsTransferringOwnership);
  
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          assertEquals(partitions(tp(topic1, 0)), assignment.get(consumer1));
          assertEquals(partitions(tp(topic1, 1), tp(topic2, 1)), assignment.get(consumer2));
          assertEquals(partitions(tp(topic2, 0)), assignment.get(consumer3));
  
          assertTrue(isFullyBalanced(assignment));
      }
  
      /**
       * This unit test is testing all consumers owned less than minQuota partitions situation
       */
      @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
      @EnumSource(RackConfig.class)
      public void testAllConsumersAreUnderMinQuota(RackConfig rackConfig) {
          initializeRacks(rackConfig);
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          partitionsPerTopic.put(topic1, partitionInfos(topic1, 2));
          partitionsPerTopic.put(topic2, partitionInfos(topic2, 3));
  
          List<String> subscribedTopics = topics(topic1, topic2);
  
          subscriptions.put(consumer1,
              buildSubscriptionV2Above(subscribedTopics, partitions(tp(topic1, 0)), generationId, 0));
          subscriptions.put(consumer2,
              buildSubscriptionV2Above(subscribedTopics, partitions(tp(topic1, 1)), generationId, 1));
          subscriptions.put(consumer3, buildSubscriptionV2Above(subscribedTopics, Collections.emptyList(), generationId, 2));
  
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
  
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          assertEquals(partitions(tp(topic1, 0), tp(topic2, 1)), assignment.get(consumer1));
          assertEquals(partitions(tp(topic1, 1), tp(topic2, 2)), assignment.get(consumer2));
          assertEquals(partitions(tp(topic2, 0)), assignment.get(consumer3));
  
          assertTrue(isFullyBalanced(assignment));
      }
  
      @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
      @EnumSource(RackConfig.class)
      public void testAddRemoveConsumerOneTopic(RackConfig rackConfig) {
          initializeRacks(rackConfig);
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          partitionsPerTopic.put(topic, partitionInfos(topic, 3));
          subscriptions.put(consumer1, subscription(topics(topic), 0));
  
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          assertEquals(partitions(tp(topic, 0), tp(topic, 1), tp(topic, 2)), assignment.get(consumer1));
  
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          assertTrue(isFullyBalanced(assignment));
  
          subscriptions.put(consumer1, buildSubscriptionV2Above(topics(topic), assignment.get(consumer1), generationId, 0));
          subscriptions.put(consumer2, buildSubscriptionV2Above(topics(topic), Collections.emptyList(), generationId, 1));
          assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          assertEquals(Collections.singletonMap(tp(topic, 2), consumer2), assignor.partitionsTransferringOwnership);
  
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          assertEquals(partitions(tp(topic, 0), tp(topic, 1)), assignment.get(consumer1));
          assertEquals(partitions(tp(topic, 2)), assignment.get(consumer2));
          assertTrue(isFullyBalanced(assignment));
  
          subscriptions.remove(consumer1);
          subscriptions.put(consumer2, buildSubscriptionV2Above(topics(topic), assignment.get(consumer2), generationId, 1));
          assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          assertEquals(new HashSet<>(partitions(tp(topic, 2), tp(topic, 1), tp(topic, 0))),
              new HashSet<>(assignment.get(consumer2)));
          assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
  
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          assertTrue(isFullyBalanced(assignment));
      }
  
      @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
      @EnumSource(RackConfig.class)
      public void testTopicBalanceAfterReassignment(RackConfig rackConfig) {
          initializeRacks(rackConfig);
          List<String> allTopics = topics(topic1, topic2);
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          partitionsPerTopic.put(topic1, partitionInfos(topic1, 12));
          partitionsPerTopic.put(topic2, partitionInfos(topic2, 12));
          subscriptions.put(consumer1, subscription(allTopics, 0));
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
  
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          assignment.forEach((consumer, tps) -> assertEquals(12, tps.stream().filter(tp -> tp.topic().equals(topic1)).count()));
          assignment.forEach((consumer, tps) -> assertEquals(12, tps.stream().filter(tp -> tp.topic().equals(topic2)).count()));
          assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
          assertTrue(isFullyBalanced(assignment));
  
          // Add another consumer
          subscriptions.put(consumer1, buildSubscriptionV2Above(allTopics, assignment.get(consumer1), generationId, 0));
          subscriptions.put(consumer2, buildSubscriptionV2Above(allTopics, Collections.emptyList(), generationId, 1));
          assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
  
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          assignment.forEach((consumer, tps) -> assertEquals(6, tps.stream().filter(tp -> tp.topic().equals(topic1)).count()));
          assignment.forEach((consumer, tps) -> assertEquals(6, tps.stream().filter(tp -> tp.topic().equals(topic2)).count()));
          assertTrue(isFullyBalanced(assignment));
  
          // Add two more consumers
          subscriptions.put(consumer1, buildSubscriptionV2Above(allTopics, assignment.get(consumer1), generationId, 0));
          subscriptions.put(consumer2, buildSubscriptionV2Above(allTopics, assignment.get(consumer2), generationId, 1));
          subscriptions.put(consumer3, buildSubscriptionV2Above(allTopics, Collections.emptyList(), generationId, 2));
          subscriptions.put(consumer4, buildSubscriptionV2Above(allTopics, Collections.emptyList(), generationId, 3));
          assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
  
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          assignment.forEach((consumer, tps) -> assertEquals(3, tps.stream().filter(tp -> tp.topic().equals(topic1)).count()));
          assignment.forEach((consumer, tps) -> assertEquals(3, tps.stream().filter(tp -> tp.topic().equals(topic2)).count()));
          assertTrue(isFullyBalanced(assignment));
  
          // remove 2 consumers
          subscriptions.remove(consumer1);
          subscriptions.remove(consumer2);
          subscriptions.put(consumer3, buildSubscriptionV2Above(allTopics, assignment.get(consumer3), generationId, 2));
          subscriptions.put(consumer4, buildSubscriptionV2Above(allTopics, assignment.get(consumer4), generationId, 3));
          assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
  
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          assignment.forEach((consumer, tps) -> assertEquals(6, tps.stream().filter(tp -> tp.topic().equals(topic1)).count()));
          assignment.forEach((consumer, tps) -> assertEquals(6, tps.stream().filter(tp -> tp.topic().equals(topic2)).count()));
          assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
          assertTrue(isFullyBalanced(assignment));
      }
  
      @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
      @EnumSource(RackConfig.class)
      public void testAddRemoveTwoConsumersTwoTopics(RackConfig rackConfig) {
          initializeRacks(rackConfig);
          List<String> allTopics = topics(topic1, topic2);
  
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          partitionsPerTopic.put(topic1, partitionInfos(topic1, 3));
          partitionsPerTopic.put(topic2, partitionInfos(topic2, 4));
          subscriptions.put(consumer1, subscription(allTopics, 0));
          subscriptions.put(consumer2, subscription(allTopics, 1));
  
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          assertEquals(partitions(tp(topic1, 0), tp(topic1, 2), tp(topic2, 1), tp(topic2, 3)), assignment.get(consumer1));
          assertEquals(partitions(tp(topic1, 1), tp(topic2, 0), tp(topic2, 2)), assignment.get(consumer2));
          assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
  
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          assertTrue(isFullyBalanced(assignment));
  
          // add 2 consumers
          subscriptions.put(consumer1, buildSubscriptionV2Above(allTopics, assignment.get(consumer1), generationId, 0));
          subscriptions.put(consumer2, buildSubscriptionV2Above(allTopics, assignment.get(consumer2), generationId, 1));
          subscriptions.put(consumer3, buildSubscriptionV2Above(allTopics, Collections.emptyList(), generationId, 2));
          subscriptions.put(consumer4, buildSubscriptionV2Above(allTopics, Collections.emptyList(), generationId, 3));
          assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
  
          Map<TopicPartition, String> expectedPartitionsTransferringOwnership = new HashMap<>();
          expectedPartitionsTransferringOwnership.put(tp(topic1, 2), consumer3);
          expectedPartitionsTransferringOwnership.put(tp(topic2, 3), consumer3);
          expectedPartitionsTransferringOwnership.put(tp(topic2, 2), consumer4);
          assertEquals(expectedPartitionsTransferringOwnership, assignor.partitionsTransferringOwnership);
  
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          assertEquals(partitions(tp(topic1, 0), tp(topic2, 1)), assignment.get(consumer1));
          assertEquals(partitions(tp(topic2, 0), tp(topic1, 1)), assignment.get(consumer2));
          assertEquals(partitions(tp(topic1, 2), tp(topic2, 3)), assignment.get(consumer3));
          assertEquals(partitions(tp(topic2, 2)), assignment.get(consumer4));
          assertTrue(isFullyBalanced(assignment));
  
          // remove 2 consumers
          subscriptions.remove(consumer1);
          subscriptions.remove(consumer2);
          subscriptions.put(consumer3, buildSubscriptionV2Above(allTopics, assignment.get(consumer3), generationId, 2));
          subscriptions.put(consumer4, buildSubscriptionV2Above(allTopics, assignment.get(consumer4), generationId, 3));
          assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          assertEquals(partitions(tp(topic1, 2), tp(topic2, 3), tp(topic1, 0), tp(topic2, 1)), assignment.get(consumer3));
          assertEquals(partitions(tp(topic2, 2), tp(topic1, 1), tp(topic2, 0)), assignment.get(consumer4));
  
          assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
  
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          assertTrue(isFullyBalanced(assignment));
      }
  
      /**
       * 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
       */
      @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
      @EnumSource(RackConfig.class)
      public void testPoorRoundRobinAssignmentScenario(RackConfig rackConfig) {
          initializeRacks(rackConfig);
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          for (int i = 1; i <= 5; i++) {
              String topicName = String.format("topic%d", i);
              partitionsPerTopic.put(topicName, partitionInfos(topicName, (i % 2) + 1));
          }
  
          subscriptions.put("consumer1",
              subscription(topics("topic1", "topic2", "topic3", "topic4", "topic5"), 0));
          subscriptions.put("consumer2",
              subscription(topics("topic1", "topic3", "topic5"), 1));
          subscriptions.put("consumer3",
              subscription(topics("topic1", "topic3", "topic5"), 2));
          subscriptions.put("consumer4",
              subscription(topics("topic1", "topic2", "topic3", "topic4", "topic5"), 3));
  
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
      }
  
      @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
      @EnumSource(RackConfig.class)
      public void testAddRemoveTopicTwoConsumers(RackConfig rackConfig) {
          initializeRacks(rackConfig);
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          partitionsPerTopic.put(topic, partitionInfos(topic, 3));
          subscriptions.put(consumer1, subscription(topics(topic), 0));
          subscriptions.put(consumer2, subscription(topics(topic), 1));
  
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
          // verify balance
          assertTrue(isFullyBalanced(assignment));
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          // verify stickiness
          List<TopicPartition> consumer1Assignment1 = assignment.get(consumer1);
          List<TopicPartition> consumer2Assignment1 = assignment.get(consumer2);
          assertTrue((consumer1Assignment1.size() == 1 && consumer2Assignment1.size() == 2) ||
              (consumer1Assignment1.size() == 2 && consumer2Assignment1.size() == 1));
  
          partitionsPerTopic.put(topic2, partitionInfos(topic2, 3));
          subscriptions.put(consumer1, buildSubscriptionV2Above(topics(topic, topic2), assignment.get(consumer1), generationId, 0));
          subscriptions.put(consumer2, buildSubscriptionV2Above(topics(topic, topic2), assignment.get(consumer2), generationId, 1));
  
          assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
          // verify balance
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          assertTrue(isFullyBalanced(assignment));
          // verify stickiness
          List<TopicPartition> consumer1assignment = assignment.get(consumer1);
          List<TopicPartition> consumer2assignment = assignment.get(consumer2);
          assertTrue(consumer1assignment.size() == 3 && consumer2assignment.size() == 3);
          assertTrue(consumer1assignment.containsAll(consumer1Assignment1));
          assertTrue(consumer2assignment.containsAll(consumer2Assignment1));
  
          partitionsPerTopic.remove(topic);
          subscriptions.put(consumer1, buildSubscriptionV2Above(topics(topic2), assignment.get(consumer1), generationId, 0));
          subscriptions.put(consumer2, buildSubscriptionV2Above(topics(topic2), assignment.get(consumer2), generationId, 1));
  
          assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
          // verify balance
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          assertTrue(isFullyBalanced(assignment));
          // verify stickiness
          List<TopicPartition> consumer1Assignment3 = assignment.get(consumer1);
          List<TopicPartition> consumer2Assignment3 = assignment.get(consumer2);
          assertTrue((consumer1Assignment3.size() == 1 && consumer2Assignment3.size() == 2) ||
              (consumer1Assignment3.size() == 2 && consumer2Assignment3.size() == 1));
          assertTrue(consumer1assignment.containsAll(consumer1Assignment3));
          assertTrue(consumer2assignment.containsAll(consumer2Assignment3));
      }
  
      @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
      @EnumSource(RackConfig.class)
      public void testReassignmentAfterOneConsumerLeaves(RackConfig rackConfig) {
          initializeRacks(rackConfig);
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          for (int i = 1; i < 20; i++) {
              String topicName = getTopicName(i, 20);
              partitionsPerTopic.put(topicName, partitionInfos(topicName, i));
          }
  
          for (int i = 1; i < 20; i++) {
              List<String> topics = new ArrayList<>();
              for (int j = 1; j <= i; j++)
                  topics.add(getTopicName(j, 20));
              subscriptions.put(getConsumerName(i, 20), subscription(topics, i));
          }
  
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
  
          for (int i = 1; i < 20; i++) {
              String consumer = getConsumerName(i, 20);
              subscriptions.put(consumer,
                  buildSubscriptionV2Above(subscriptions.get(consumer).topics(), assignment.get(consumer), generationId, i));
          }
          subscriptions.remove("consumer10");
  
          assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          assertTrue(assignor.isSticky());
      }
  
  
      @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
      @EnumSource(RackConfig.class)
      public void testReassignmentAfterOneConsumerAdded(RackConfig rackConfig) {
          initializeRacks(rackConfig);
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          partitionsPerTopic.put("topic", partitionInfos("topic", 20));
  
          for (int i = 1; i < 10; i++)
              subscriptions.put(getConsumerName(i, 10),
                  subscription(topics("topic"), i));
  
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
  
          // add a new consumer
          subscriptions.put(getConsumerName(10, 10), subscription(topics("topic"), 10));
  
          assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
      }
  
      @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
      @EnumSource(RackConfig.class)
      public void testSameSubscriptions(RackConfig rackConfig) {
          initializeRacks(rackConfig);
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          for (int i = 1; i < 15; i++) {
              String topicName = getTopicName(i, 15);
              partitionsPerTopic.put(topicName, partitionInfos(topicName, i));
          }
  
          for (int i = 1; i < 9; i++) {
              List<String> topics = new ArrayList<>();
              for (int j = 1; j <= partitionsPerTopic.size(); j++)
                  topics.add(getTopicName(j, 15));
              subscriptions.put(getConsumerName(i, 9), subscription(topics, i));
          }
  
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
  
          for (int i = 1; i < 9; i++) {
              String consumer = getConsumerName(i, 9);
              subscriptions.put(consumer,
                  buildSubscriptionV2Above(subscriptions.get(consumer).topics(), assignment.get(consumer), generationId, i));
          }
          subscriptions.remove(getConsumerName(5, 9));
  
          assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
      }
  
      @Timeout(30)
      @ParameterizedTest(name = TEST_NAME_WITH_CONSUMER_RACK)
      @ValueSource(booleans = {false, true})
      public void testLargeAssignmentAndGroupWithUniformSubscription(boolean hasConsumerRack) {
          initializeRacks(hasConsumerRack ? RackConfig.BROKER_AND_CONSUMER_RACK : RackConfig.NO_CONSUMER_RACK);
          // 1 million partitions for non-rack-aware! For rack-aware, use smaller number of partitions to reduce test run time.
          int topicCount = hasConsumerRack ? 50 : 500;
          int partitionCount = 2_000;
          int consumerCount = 2_000;
  
          List<String> topics = new ArrayList<>();
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          for (int i = 0; i < topicCount; i++) {
              String topicName = getTopicName(i, topicCount);
              topics.add(topicName);
              partitionsPerTopic.put(topicName, partitionInfos(topicName, partitionCount));
          }
  
          for (int i = 0; i < consumerCount; i++) {
              subscriptions.put(getConsumerName(i, consumerCount), subscription(topics, i));
          }
  
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
  
          for (int i = 1; i < consumerCount; i++) {
              String consumer = getConsumerName(i, consumerCount);
              subscriptions.put(consumer, buildSubscriptionV2Above(topics, assignment.get(consumer), generationId, i));
          }
  
          assignor.assignPartitions(partitionsPerTopic, subscriptions);
      }
  
      @Flaky("KAFKA-13514")
      @Timeout(90)
      @ParameterizedTest(name = TEST_NAME_WITH_CONSUMER_RACK)
      @ValueSource(booleans = {false, true})
      public void testLargeAssignmentAndGroupWithNonEqualSubscription(boolean hasConsumerRack) {
          initializeRacks(hasConsumerRack ? RackConfig.BROKER_AND_CONSUMER_RACK : RackConfig.NO_CONSUMER_RACK);
          // 1 million partitions for non-rack-aware! For rack-aware, use smaller number of partitions to reduce test run time.
          int topicCount = hasConsumerRack ? 50 : 500;
          int partitionCount = 2_000;
          int consumerCount = 2_000;
  
          List<String> topics = new ArrayList<>();
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          for (int i = 0; i < topicCount; i++) {
              String topicName = getTopicName(i, topicCount);
              topics.add(topicName);
              partitionsPerTopic.put(topicName, partitionInfos(topicName, partitionCount));
          }
          for (int i = 0; i < consumerCount; i++) {
              if (i == consumerCount - 1) {
                  subscriptions.put(getConsumerName(i, consumerCount), subscription(topics.subList(0, 1), i));
              } else {
                  subscriptions.put(getConsumerName(i, consumerCount), subscription(topics, i));
              }
          }
  
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
  
          for (int i = 1; i < consumerCount; i++) {
              String consumer = getConsumerName(i, consumerCount);
              if (i == consumerCount - 1) {
                  subscriptions.put(consumer, buildSubscriptionV2Above(topics.subList(0, 1), assignment.get(consumer), generationId, i));
              } else {
                  subscriptions.put(consumer, buildSubscriptionV2Above(topics, assignment.get(consumer), generationId, i));
              }
          }
  
          assignor.assignPartitions(partitionsPerTopic, subscriptions);
      }
  
      @Timeout(90)
      @ParameterizedTest(name = TEST_NAME_WITH_CONSUMER_RACK)
      @ValueSource(booleans = {false, true})
      public void testAssignmentAndGroupWithNonEqualSubscriptionNotTimeout(boolean hasConsumerRack) {
          initializeRacks(hasConsumerRack ? RackConfig.BROKER_AND_CONSUMER_RACK : RackConfig.NO_CONSUMER_RACK);
          int topicCount = hasConsumerRack ? 50 : 100;
          int partitionCount = 2_00;
          int consumerCount = 5_00;
  
          List<String> topics = new ArrayList<>();
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          for (int i = 0; i < topicCount; i++) {
              String topicName = getTopicName(i, topicCount);
              topics.add(topicName);
              partitionsPerTopic.put(topicName, partitionInfos(topicName, partitionCount));
          }
          for (int i = 0; i < consumerCount; i++) {
              if (i % 4 == 0) {
                  subscriptions.put(getConsumerName(i, consumerCount),
                          subscription(topics.subList(0, topicCount / 2), i));
              } else {
                  subscriptions.put(getConsumerName(i, consumerCount),
                          subscription(topics.subList(topicCount / 2, topicCount), i));
              }
          }
  
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
  
          for (int i = 1; i < consumerCount; i++) {
              String consumer = getConsumerName(i, consumerCount);
              if (i % 4 == 0) {
                  subscriptions.put(
                          consumer,
                          buildSubscriptionV2Above(topics.subList(0, topicCount / 2),
                          assignment.get(consumer), generationId, i)
                  );
              } else {
                  subscriptions.put(
                          consumer,
                          buildSubscriptionV2Above(topics.subList(topicCount / 2, topicCount),
                          assignment.get(consumer), generationId, i)
                  );
              }
          }
  
          assignor.assignPartitions(partitionsPerTopic, subscriptions);
      }
  
      @Test
      public void testSubscriptionNotEqualAndAssignSamePartitionWith3Generation() {
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          partitionsPerTopic.put(topic, partitionInfos(topic, 6));
          partitionsPerTopic.put(topic1, partitionInfos(topic1, 1));
          int[][] sequence = new int[][]{{1, 2, 3}, {1, 3, 2}, {2, 1, 3}, {2, 3, 1}, {3, 1, 2}, {3, 2, 1}};
          for (int[] ints : sequence) {
              subscriptions.put(
                      consumer1,
                      buildSubscriptionV2Above(topics(topic),
                      partitions(tp(topic, 0), tp(topic, 2)), ints[0], 0)
              );
              subscriptions.put(
                      consumer2,
                      buildSubscriptionV2Above(topics(topic),
                      partitions(tp(topic, 1), tp(topic, 2), tp(topic, 3)), ints[1], 1)
              );
              subscriptions.put(
                      consumer3,
                      buildSubscriptionV2Above(topics(topic),
                      partitions(tp(topic, 2), tp(topic, 4), tp(topic, 5)), ints[2], 2)
              );
              subscriptions.put(
                      consumer4,
                      buildSubscriptionV2Above(topics(topic1),
                      partitions(tp(topic1, 0)), 2, 3)
              );
  
              Map<String, List<TopicPartition>> assign = assignor.assignPartitions(partitionsPerTopic, subscriptions);
              assertEquals(assign.values().stream().mapToInt(List::size).sum(),
                      assign.values().stream().flatMap(List::stream).collect(Collectors.toSet()).size());
              for (List<TopicPartition> list: assign.values()) {
                  assertTrue(list.size() >= 1 && list.size() <= 2);
              }
          }
      }
  
      @Timeout(60)
      @ParameterizedTest(name = TEST_NAME_WITH_CONSUMER_RACK)
      @ValueSource(booleans = {false, true})
      public void testLargeAssignmentWithMultipleConsumersLeavingAndRandomSubscription(boolean hasConsumerRack) {
          initializeRacks(hasConsumerRack ? RackConfig.BROKER_AND_CONSUMER_RACK : RackConfig.NO_CONSUMER_RACK);
          Random rand = new Random();
          int topicCount = 40;
          int consumerCount = 200;
  
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          for (int i = 0; i < topicCount; i++) {
              String topicName = getTopicName(i, topicCount);
              partitionsPerTopic.put(topicName, partitionInfos(topicName, rand.nextInt(10) + 1));
          }
  
          for (int i = 0; i < consumerCount; i++) {
              List<String> topics = new ArrayList<>();
              for (int j = 0; j < rand.nextInt(20); j++)
                  topics.add(getTopicName(rand.nextInt(topicCount), topicCount));
              subscriptions.put(getConsumerName(i, consumerCount), subscription(topics, i));
          }
  
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
  
          for (int i = 1; i < consumerCount; i++) {
              String consumer = getConsumerName(i, consumerCount);
              subscriptions.put(consumer,
                  buildSubscriptionV2Above(subscriptions.get(consumer).topics(), assignment.get(consumer), generationId, i));
          }
          for (int i = 0; i < 50; ++i) {
              String c = getConsumerName(rand.nextInt(consumerCount), consumerCount);
              subscriptions.remove(c);
          }
  
          assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          assertTrue(assignor.isSticky());
      }
  
      @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
      @EnumSource(RackConfig.class)
      public void testNewSubscription(RackConfig rackConfig) {
          initializeRacks(rackConfig);
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          for (int i = 1; i < 5; i++) {
              String topicName = getTopicName(i, 5);
              partitionsPerTopic.put(topicName, partitionInfos(topicName, 1));
          }
  
          for (int i = 0; i < 3; i++) {
              List<String> topics = new ArrayList<>();
              for (int j = i; j <= 3 * i - 2; j++)
                  topics.add(getTopicName(j, 5));
              subscriptions.put(getConsumerName(i, 3), subscription(topics, i));
          }
  
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
  
          subscriptions.get(getConsumerName(0, 3)).topics().add(getTopicName(1, 5));
  
          assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          assertTrue(assignor.isSticky());
      }
  
      @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
      @EnumSource(RackConfig.class)
      public void testMoveExistingAssignments(RackConfig rackConfig) {
          initializeRacks(rackConfig);
          String topic4 = "topic4";
          String topic5 = "topic5";
          String topic6 = "topic6";
  
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          for (int i = 1; i <= 6; i++) {
              String topicName = String.format("topic%d", i);
              partitionsPerTopic.put(topicName, partitionInfos(topicName, 1));
          }
  
          subscriptions.put(consumer1,
              buildSubscriptionV2Above(topics(topic1, topic2),
                  partitions(tp(topic1, 0)), generationId, 0));
          subscriptions.put(consumer2,
              buildSubscriptionV2Above(topics(topic1, topic2, topic3, topic4),
                  partitions(tp(topic2, 0), tp(topic3, 0)), generationId, 1));
          subscriptions.put(consumer3,
              buildSubscriptionV2Above(topics(topic2, topic3, topic4, topic5, topic6),
                  partitions(tp(topic4, 0), tp(topic5, 0), tp(topic6, 0)), generationId, 2));
  
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          assertNull(assignor.partitionsTransferringOwnership);
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
      }
  
      @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);
          }
      }
  
      @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
      @EnumSource(RackConfig.class)
      public void testAssignmentUpdatedForDeletedTopic(RackConfig rackConfig) {
          initializeRacks(rackConfig);
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          partitionsPerTopic.put(topic1, partitionInfos(topic1, 1));
          partitionsPerTopic.put(topic3, partitionInfos(topic3, 100));
          subscriptions = Collections.singletonMap(consumerId, subscription(topics(topic1, topic2, topic3), 0));
  
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
          assertEquals(assignment.values().stream().mapToInt(List::size).sum(), 1 + 100);
          assertEquals(Collections.singleton(consumerId), assignment.keySet());
          assertTrue(isFullyBalanced(assignment));
      }
  
      @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
      @EnumSource(RackConfig.class)
      public void testNoExceptionThrownWhenOnlySubscribedTopicDeleted(RackConfig rackConfig) {
          initializeRacks(rackConfig);
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          partitionsPerTopic.put(topic, partitionInfos(topic, 3));
          subscriptions.put(consumerId, subscription(topics(topic), 0));
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
          subscriptions.put(consumerId, buildSubscriptionV2Above(topics(topic), assignment.get(consumerId), generationId, 0));
  
          assignment = assignor.assign(Collections.emptyMap(), subscriptions);
          assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
          assertEquals(assignment.size(), 1);
          assertTrue(assignment.get(consumerId).isEmpty());
      }
  
      @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
      @EnumSource(RackConfig.class)
      public void testReassignmentWithRandomSubscriptionsAndChanges(RackConfig rackConfig) {
          initializeRacks(rackConfig);
          final int minNumConsumers = 20;
          final int maxNumConsumers = 40;
          final int minNumTopics = 10;
          final int maxNumTopics = 20;
  
          for (int round = 1; round <= 100; ++round) {
              int numTopics = minNumTopics + new Random().nextInt(maxNumTopics - minNumTopics);
  
              ArrayList<String> topics = new ArrayList<>();
  
              Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
              for (int i = 0; i < numTopics; ++i) {
                  String topicName = getTopicName(i, maxNumTopics);
                  topics.add(topicName);
                  partitionsPerTopic.put(topicName, partitionInfos(topicName, i + 1));
              }
  
              int numConsumers = minNumConsumers + new Random().nextInt(maxNumConsumers - minNumConsumers);
  
              for (int i = 0; i < numConsumers; ++i) {
                  List<String> sub = Utils.sorted(getRandomSublist(topics));
                  subscriptions.put(getConsumerName(i, maxNumConsumers), subscription(sub, i));
              }
  
              assignor = createAssignor();
  
              Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
              verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
  
              subscriptions.clear();
              for (int i = 0; i < numConsumers; ++i) {
                  List<String> sub = Utils.sorted(getRandomSublist(topics));
                  String consumer = getConsumerName(i, maxNumConsumers);
                  subscriptions.put(consumer, buildSubscriptionV2Above(sub, assignment.get(consumer), generationId, i));
              }
  
              assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
              verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
              assertTrue(assignor.isSticky());
          }
      }
  
      @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
      @EnumSource(RackConfig.class)
      public void testAllConsumersReachExpectedQuotaAndAreConsideredFilled(RackConfig rackConfig) {
          initializeRacks(rackConfig);
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          partitionsPerTopic.put(topic, partitionInfos(topic, 4));
  
          subscriptions.put(consumer1, buildSubscriptionV2Above(topics(topic), partitions(tp(topic, 0), tp(topic, 1)), generationId, 0));
          subscriptions.put(consumer2, buildSubscriptionV2Above(topics(topic), partitions(tp(topic, 2)), generationId, 1));
          subscriptions.put(consumer3, buildSubscriptionV2Above(topics(topic), Collections.emptyList(), generationId, 2));
  
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          assertEquals(partitions(tp(topic, 0), tp(topic, 1)), assignment.get(consumer1));
          assertEquals(partitions(tp(topic, 2)), assignment.get(consumer2));
          assertEquals(partitions(tp(topic, 3)), assignment.get(consumer3));
          assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
  
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          assertTrue(isFullyBalanced(assignment));
      }
  
      @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
      @EnumSource(RackConfig.class)
      public void testOwnedPartitionsAreInvalidatedForConsumerWithStaleGeneration(RackConfig rackConfig) {
          initializeRacks(rackConfig);
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          partitionsPerTopic.put(topic, partitionInfos(topic, 3));
          partitionsPerTopic.put(topic2, partitionInfos(topic2, 3));
  
          int currentGeneration = 10;
  
          subscriptions.put(consumer1, buildSubscriptionV2Above(topics(topic, topic2), partitions(tp(topic, 0), tp(topic, 2), tp(topic2, 1)), currentGeneration, 0));
          subscriptions.put(consumer2, buildSubscriptionV2Above(topics(topic, topic2), partitions(tp(topic, 0), tp(topic, 2), tp(topic2, 1)), currentGeneration - 1, 1));
  
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          assertEquals(new HashSet<>(partitions(tp(topic, 0), tp(topic, 2), tp(topic2, 1))), new HashSet<>(assignment.get(consumer1)));
          assertEquals(new HashSet<>(partitions(tp(topic, 1), tp(topic2, 0), tp(topic2, 2))), new HashSet<>(assignment.get(consumer2)));
          assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
  
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          assertTrue(isFullyBalanced(assignment));
      }
  
      @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
      @EnumSource(RackConfig.class)
      public void testOwnedPartitionsAreInvalidatedForConsumerWithNoGeneration(RackConfig rackConfig) {
          initializeRacks(rackConfig);
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          partitionsPerTopic.put(topic, partitionInfos(topic, 3));
          partitionsPerTopic.put(topic2, partitionInfos(topic2, 3));
  
          int currentGeneration = 10;
  
          subscriptions.put(consumer1, buildSubscriptionV2Above(topics(topic, topic2), partitions(tp(topic, 0), tp(topic, 2), tp(topic2, 1)), currentGeneration, 0));
          subscriptions.put(consumer2, buildSubscriptionV2Above(topics(topic, topic2), partitions(tp(topic, 0), tp(topic, 2), tp(topic2, 1)), DEFAULT_GENERATION, 1));
  
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          assertEquals(new HashSet<>(partitions(tp(topic, 0), tp(topic, 2), tp(topic2, 1))), new HashSet<>(assignment.get(consumer1)));
          assertEquals(new HashSet<>(partitions(tp(topic, 1), tp(topic2, 0), tp(topic2, 2))), new HashSet<>(assignment.get(consumer2)));
          assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
  
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          assertTrue(isFullyBalanced(assignment));
      }
  
      @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
      @EnumSource(RackConfig.class)
      public void testPartitionsTransferringOwnershipIncludeThePartitionClaimedByMultipleConsumersInSameGeneration(RackConfig rackConfig) {
          initializeRacks(rackConfig);
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          partitionsPerTopic.put(topic, partitionInfos(topic, 3));
  
          // partition topic-0 is owned by multiple consumer
          subscriptions.put(consumer1, buildSubscriptionV2Above(topics(topic), partitions(tp(topic, 0), tp(topic, 1)), generationId, 0));
          subscriptions.put(consumer2, buildSubscriptionV2Above(topics(topic), partitions(tp(topic, 0), tp(topic, 2)), generationId, 1));
          subscriptions.put(consumer3, buildSubscriptionV2Above(topics(topic), emptyList(), generationId, 2));
  
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          // we should include the partitions claimed by multiple consumers in partitionsTransferringOwnership
          assertEquals(Collections.singletonMap(tp(topic, 0), consumer3), assignor.partitionsTransferringOwnership);
  
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          assertEquals(partitions(tp(topic, 1)), assignment.get(consumer1));
          assertEquals(partitions(tp(topic, 2)), assignment.get(consumer2));
          assertEquals(partitions(tp(topic, 0)), assignment.get(consumer3));
          assertTrue(isFullyBalanced(assignment));
      }
  
      @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
      @EnumSource(RackConfig.class)
      public void testEnsurePartitionsAssignedToHighestGeneration(RackConfig rackConfig) {
          initializeRacks(rackConfig);
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          partitionsPerTopic.put(topic, partitionInfos(topic, 3));
          partitionsPerTopic.put(topic2, partitionInfos(topic2, 3));
          partitionsPerTopic.put(topic3, partitionInfos(topic3, 3));
  
          int currentGeneration = 10;
  
          // ensure partitions are always assigned to the member with the highest generation
          subscriptions.put(consumer1, buildSubscriptionV2Above(topics(topic, topic2, topic3),
              partitions(tp(topic, 0), tp(topic2, 0), tp(topic3, 0)), currentGeneration, 0));
          subscriptions.put(consumer2, buildSubscriptionV2Above(topics(topic, topic2, topic3),
              partitions(tp(topic, 1), tp(topic2, 1), tp(topic3, 1)), currentGeneration - 1, 1));
          subscriptions.put(consumer3, buildSubscriptionV2Above(topics(topic, topic2, topic3),
              partitions(tp(topic2, 1), tp(topic3, 0), tp(topic3, 2)), currentGeneration - 2, 1));
  
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          assertEquals(new HashSet<>(partitions(tp(topic, 0), tp(topic2, 0), tp(topic3, 0))),
              new HashSet<>(assignment.get(consumer1)));
          assertEquals(new HashSet<>(partitions(tp(topic, 1), tp(topic2, 1), tp(topic3, 1))),
              new HashSet<>(assignment.get(consumer2)));
          assertEquals(new HashSet<>(partitions(tp(topic, 2), tp(topic2, 2), tp(topic3, 2))),
              new HashSet<>(assignment.get(consumer3)));
          assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
  
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          assertTrue(isFullyBalanced(assignment));
      }
  
      @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
      @EnumSource(RackConfig.class)
      public void testNoReassignmentOnCurrentMembers(RackConfig rackConfig) {
          initializeRacks(rackConfig);
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          partitionsPerTopic.put(topic, partitionInfos(topic, 3));
          partitionsPerTopic.put(topic1, partitionInfos(topic1, 3));
          partitionsPerTopic.put(topic2, partitionInfos(topic2, 3));
          partitionsPerTopic.put(topic3, partitionInfos(topic3, 3));
  
          int currentGeneration = 10;
  
          subscriptions.put(consumer1, buildSubscriptionV2Above(topics(topic, topic2, topic3, topic1),
              partitions(), DEFAULT_GENERATION, 0));
          subscriptions.put(consumer2, buildSubscriptionV2Above(topics(topic, topic2, topic3, topic1),
              partitions(tp(topic, 0), tp(topic2, 0), tp(topic1, 0)), currentGeneration - 1, 1));
          subscriptions.put(consumer3, buildSubscriptionV2Above(topics(topic, topic2, topic3, topic1),
              partitions(tp(topic3, 2), tp(topic2, 2), tp(topic1, 1)), currentGeneration - 2, 2));
          subscriptions.put(consumer4, buildSubscriptionV2Above(topics(topic, topic2, topic3, topic1),
              partitions(tp(topic3, 1), tp(topic, 1), tp(topic, 2)), currentGeneration - 3, 3));
  
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          // ensure assigned partitions don't get reassigned
          assertEquals(new HashSet<>(partitions(tp(topic1, 2), tp(topic2, 1), tp(topic3, 0))),
                  new HashSet<>(assignment.get(consumer1)));
          assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
  
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          assertTrue(isFullyBalanced(assignment));
      }
  
      @ParameterizedTest(name = TEST_NAME_WITH_RACK_CONFIG)
      @EnumSource(RackConfig.class)
      public void testOwnedPartitionsAreInvalidatedForConsumerWithMultipleGeneration(RackConfig rackConfig) {
          initializeRacks(rackConfig);
          Map<String, List<PartitionInfo>> partitionsPerTopic = new HashMap<>();
          partitionsPerTopic.put(topic, partitionInfos(topic, 3));
          partitionsPerTopic.put(topic2, partitionInfos(topic2, 3));
  
          int currentGeneration = 10;
  
          subscriptions.put(consumer1, buildSubscriptionV2Above(topics(topic, topic2),
              partitions(tp(topic, 0), tp(topic2, 1), tp(topic, 1)), currentGeneration, 0));
          subscriptions.put(consumer2, buildSubscriptionV2Above(topics(topic, topic2),
              partitions(tp(topic, 0), tp(topic2, 1), tp(topic2, 2)), currentGeneration - 2, 1));
  
          Map<String, List<TopicPartition>> assignment = assignor.assignPartitions(partitionsPerTopic, subscriptions);
          assertEquals(new HashSet<>(partitions(tp(topic, 0), tp(topic2, 1), tp(topic, 1))),
              new HashSet<>(assignment.get(consumer1)));
          assertEquals(new HashSet<>(partitions(tp(topic, 2), tp(topic2, 2), tp(topic2, 0))),
              new HashSet<>(assignment.get(consumer2)));
          assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
  
          verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
          assertTrue(isFullyBalanced(assignment));
      }
  
      @Test
      public void testRackAwareAssignmentWithUniformSubscription() {
          Map<String, Integer> topics = mkMap(mkEntry("t1", 6), mkEntry("t2", 7), mkEntry("t3", 2));
          List<String> allTopics = asList("t1", "t2", "t3");
          List<List<String>> consumerTopics = asList(allTopics, allTopics, allTopics);
          List<String> nonRackAwareAssignment = asList(
                  "t1-0, t1-3, t2-0, t2-3, t2-6",
                  "t1-1, t1-4, t2-1, t2-4, t3-0",
                  "t1-2, t1-5, t2-2, t2-5, t3-1"
          );
          verifyUniformSubscription(assignor, topics, 3, nullRacks(3), racks(3), consumerTopics, nonRackAwareAssignment, -1);
          verifyUniformSubscription(assignor, topics, 3, racks(3), nullRacks(3), consumerTopics, nonRackAwareAssignment, -1);
          AbstractPartitionAssignorTest.preferRackAwareLogic(assignor, true);
          verifyUniformSubscription(assignor, topics, 3, racks(3), racks(3), consumerTopics, nonRackAwareAssignment, 0);
          AbstractPartitionAssignorTest.preferRackAwareLogic(assignor, true);
          verifyUniformSubscription(assignor, topics, 4, racks(4), racks(3), consumerTopics, nonRackAwareAssignment, 0);
          AbstractPartitionAssignorTest.preferRackAwareLogic(assignor, false);
          verifyUniformSubscription(assignor, topics, 3, racks(3), racks(3), consumerTopics, nonRackAwareAssignment, 0);
          verifyUniformSubscription(assignor, topics, 3, racks(3), asList("d", "e", "f"), consumerTopics, nonRackAwareAssignment, -1);
          verifyUniformSubscription(assignor, topics, 3, racks(3), asList(null, "e", "f"), consumerTopics, nonRackAwareAssignment, -1);
  
          // Verify assignment is rack-aligned for lower replication factor where brokers have a subset of partitions
          List<String> assignment = asList("t1-0, t1-3, t2-0, t2-3, t2-6", "t1-1, t1-4, t2-1, t2-4, t3-0", "t1-2, t1-5, t2-2, t2-5, t3-1");
          verifyUniformSubscription(assignor, topics, 1, racks(3), racks(3), consumerTopics, assignment, 0);
          assignment = asList("t1-0, t1-3, t2-0, t2-3, t2-6", "t1-1, t1-4, t2-1, t2-4, t3-0", "t1-2, t1-5, t2-2, t2-5, t3-1");
          verifyUniformSubscription(assignor, topics, 2, racks(3), racks(3), consumerTopics, assignment, 0);
  
          // One consumer on a rack with no partitions. We allocate with misaligned rack to this consumer to maintain balance.
          assignment = asList("t1-0, t1-3, t2-0, t2-3, t2-6", "t1-1, t1-4, t2-1, t2-4, t3-0", "t1-2, t1-5, t2-2, t2-5, t3-1");
          verifyUniformSubscription(assignor, topics, 3, racks(2), racks(3), consumerTopics, assignment, 5);
  
          if (assignor instanceof CooperativeStickyAssignor) {
              // Verify that rack-awareness is improved if already owned partitions are misaligned
              assignment = asList("t1-0, t1-3", "t2-1", "t2-5, t3-1");
              List<String> owned = asList("t1-0, t1-1, t1-2, t1-3, t1-4", "t1-5, t2-0, t2-1, t2-2, t2-3", "t2-4, t2-5, t2-6, t3-0, t3-1");
              verifyRackAssignment(assignor, topics, 1, racks(3), racks(3), consumerTopics, owned, assignment, 0);
  
              // Verify that stickiness is retained when racks match
              assignment = asList("t1-0, t1-3, t2-0, t2-3, t2-6", "t1-1, t1-4, t2-1, t2-4, t3-0", "t1-2, t1-5, t2-2, t2-5, t3-1");
              AbstractPartitionAssignorTest.preferRackAwareLogic(assignor, true);
              verifyRackAssignment(assignor, topics, 3, racks(3), racks(3), consumerTopics, assignment, assignment, 0);
          } else {
              // Verify that rack-awareness is improved if already owned partitions are misaligned
              assignment = asList("t1-0, t1-3, t2-0, t2-3, t2-6", "t1-1, t1-4, t2-1, t2-4, t3-0", "t1-2, t1-5, t2-2, t2-5, t3-1");
              List<String> owned = asList("t1-0, t1-1, t1-2, t1-3, t1-4", "t1-5, t2-0, t2-1, t2-2, t2-3", "t2-4, t2-5, t2-6, t3-0, t3-1");
              verifyRackAssignment(assignor, topics, 1, racks(3), racks(3), consumerTopics, owned, assignment, 0);
  
              // Verify that stickiness is retained when racks match
              AbstractPartitionAssignorTest.preferRackAwareLogic(assignor, true);
              verifyRackAssignment(assignor, topics, 3, racks(3), racks(3), consumerTopics, assignment, assignment, 0);
          }
      }
  
      private void verifyUniformSubscription(AbstractStickyAssignor assignor,
                                             Map<String, Integer> numPartitionsPerTopic,
                                             int replicationFactor,
                                             List<String> brokerRacks,
                                             List<String> consumerRacks,
                                             List<List<String>> consumerTopics,
                                             List<String> expectedAssignments,
                                             int numPartitionsWithRackMismatch) {
          verifyRackAssignment(assignor, numPartitionsPerTopic, replicationFactor, brokerRacks, consumerRacks,
                  consumerTopics, null, expectedAssignments, numPartitionsWithRackMismatch);
          verifyRackAssignment(assignor, numPartitionsPerTopic, replicationFactor, brokerRacks, consumerRacks,
                  consumerTopics, expectedAssignments, expectedAssignments, numPartitionsWithRackMismatch);
      }
  
      @Test
      public void testRackAwareAssignmentWithNonEqualSubscription() {
          Map<String, Integer> topics = mkMap(mkEntry("t1", 6), mkEntry("t2", 7), mkEntry("t3", 2));
          List<String> allTopics = asList("t1", "t2", "t3");
          List<List<String>> consumerTopics = asList(allTopics, allTopics, asList("t1", "t3"));
          List<String> nonRackAwareAssignment = asList(
                  "t1-5, t2-0, t2-2, t2-4, t2-6",
                  "t1-3, t2-1, t2-3, t2-5, t3-0",
                  "t1-0, t1-1, t1-2, t1-4, t3-1"
          );
          verifyNonEqualSubscription(assignor, topics, 3, nullRacks(3), racks(3), consumerTopics, nonRackAwareAssignment, -1);
          verifyNonEqualSubscription(assignor, topics, 3, racks(3), nullRacks(3), consumerTopics, nonRackAwareAssignment, -1);
          AbstractPartitionAssignorTest.preferRackAwareLogic(assignor, true);
          verifyNonEqualSubscription(assignor, topics, 3, racks(3), racks(3), consumerTopics, nonRackAwareAssignment, 0);
          AbstractPartitionAssignorTest.preferRackAwareLogic(assignor, true);
          verifyNonEqualSubscription(assignor, topics, 4, racks(4), racks(3), consumerTopics, nonRackAwareAssignment, 0);
          AbstractPartitionAssignorTest.preferRackAwareLogic(assignor, false);
          verifyNonEqualSubscription(assignor, topics, 3, racks(3), racks(3), consumerTopics, nonRackAwareAssignment, 0);
          verifyNonEqualSubscription(assignor, topics, 3, racks(3), asList("d", "e", "f"), consumerTopics, nonRackAwareAssignment, -1);
          verifyNonEqualSubscription(assignor, topics, 3, racks(3), asList(null, "e", "f"), consumerTopics, nonRackAwareAssignment, -1);
  
          // Verify assignment is rack-aligned for lower replication factor where brokers have a subset of partitions
          // Rack-alignment is best-effort, misalignments can occur when number of rack choices is low.
          List<String> assignment = asList("t1-3, t2-0, t2-2, t2-3, t2-6", "t1-4, t2-1, t2-4, t2-5, t3-0", "t1-0, t1-1, t1-2, t1-5, t3-1");
          verifyNonEqualSubscription(assignor, topics, 1, racks(3), racks(3), consumerTopics, assignment, 4);
          assignment = asList("t1-3, t2-0, t2-2, t2-5, t2-6", "t1-0, t2-1, t2-3, t2-4, t3-0", "t1-1, t1-2, t1-4, t1-5, t3-1");
          verifyNonEqualSubscription(assignor, topics, 2, racks(3), racks(3), consumerTopics, assignment, 0);
  
          // One consumer on a rack with no partitions. We allocate with misaligned rack to this consumer to maintain balance.
          verifyNonEqualSubscription(assignor, topics, 3, racks(2), racks(3), consumerTopics,
                  asList("t1-5, t2-0, t2-2, t2-4, t2-6", "t1-3, t2-1, t2-3, t2-5, t3-0", "t1-0, t1-1, t1-2, t1-4, t3-1"), 5);
  
          // Verify that rack-awareness is improved if already owned partitions are misaligned.
          // Rack alignment is attempted, but not guaranteed.
          List<String> owned = asList("t1-0, t1-1, t1-2, t1-3, t1-4", "t1-5, t2-0, t2-1, t2-2, t2-3", "t2-4, t2-5, t2-6, t3-0, t3-1");
          if (assignor instanceof StickyAssignor) {
              assignment = asList("t1-3, t2-0, t2-2, t2-3, t2-6", "t1-4, t2-1, t2-4, t2-5, t3-0", "t1-0, t1-1, t1-2, t1-5, t3-1");
              verifyRackAssignment(assignor, topics, 1, racks(3), racks(3), consumerTopics, owned, assignment, 4);
          } else {
              List<String> intermediate = asList("t1-3", "t2-1", "t3-1");
              verifyRackAssignment(assignor, topics, 1, racks(3), racks(3), consumerTopics, owned, intermediate, 0);
              assignment = asList("t1-3, t2-0, t2-2, t2-3, t2-6", "t1-4, t2-1, t2-4, t2-5, t3-0", "t1-0, t1-1, t1-2, t1-5, t3-1");
              verifyRackAssignment(assignor, topics, 1, racks(3), racks(3), consumerTopics, intermediate, assignment, 4);
          }
  
          // Verify that result is same as non-rack-aware assignment if all racks match
          if (assignor instanceof StickyAssignor) {
              assignment = asList("t1-5, t2-0, t2-2, t2-4, t2-6", "t1-3, t2-1, t2-3, t2-5, t3-0", "t1-0, t1-1, t1-2, t1-4, t3-1");
              AbstractPartitionAssignorTest.preferRackAwareLogic(assignor, false);
              verifyRackAssignment(assignor, topics, 3, racks(3), racks(3), consumerTopics, owned, assignment, 0);
              AbstractPartitionAssignorTest.preferRackAwareLogic(assignor, true);
              verifyRackAssignment(assignor, topics, 3, racks(3), racks(3), consumerTopics, owned, assignment, 0);
          } else {
              assignment = asList("t1-2, t1-3, t1-4, t2-4, t2-5", "t2-0, t2-1, t2-2, t2-3, t2-6", "t1-0, t1-1, t1-5, t3-0, t3-1");
              List<String> intermediate = asList("t1-2, t1-3, t1-4", "t2-0, t2-1, t2-2, t2-3", "t3-0, t3-1");
              AbstractPartitionAssignorTest.preferRackAwareLogic(assignor, false);
              verifyRackAssignment(assignor, topics, 3, racks(3), racks(3), consumerTopics, owned, intermediate, 0);
              verifyRackAssignment(assignor, topics, 3, racks(3), racks(3), consumerTopics, intermediate, assignment, 0);
              AbstractPartitionAssignorTest.preferRackAwareLogic(assignor, true);
              verifyRackAssignment(assignor, topics, 3, racks(3), racks(3), consumerTopics, owned, intermediate, 0);
              verifyRackAssignment(assignor, topics, 3, racks(3), racks(3), consumerTopics, intermediate, assignment, 0);
          }
      }
  
      private void verifyNonEqualSubscription(AbstractStickyAssignor assignor,
                                              Map<String, Integer> numPartitionsPerTopic,
                                              int replicationFactor,
                                              List<String> brokerRacks,
                                              List<String> consumerRacks,
                                              List<List<String>> consumerTopics,
                                              List<String> expectedAssignments,
                                              int numPartitionsWithRackMismatch) {
          verifyRackAssignment(assignor, numPartitionsPerTopic, replicationFactor, brokerRacks,
                  consumerRacks, consumerTopics, null, expectedAssignments, numPartitionsWithRackMismatch);
          verifyRackAssignment(assignor, numPartitionsPerTopic, replicationFactor, brokerRacks,
                  consumerRacks, consumerTopics, expectedAssignments, expectedAssignments, numPartitionsWithRackMismatch);
      }
  
      private String getTopicName(int i, int maxNum) {
          return getCanonicalName("t", i, maxNum);
      }
  
      private String getConsumerName(int i, int maxNum) {
          return getCanonicalName("c", i, maxNum);
      }
  
      private String getCanonicalName(String str, int i, int maxNum) {
          return str + pad(i, Integer.toString(maxNum).length());
      }
  
      private String pad(int num, int digits) {
          int iDigits = Integer.toString(num).length();
          return "0".repeat(Math.max(0, digits - iDigits)) + num;
      }
  
      protected static List<String> topics(String... topics) {
          return Arrays.asList(topics);
      }
  
      protected static List<TopicPartition> partitions(TopicPartition... partitions) {
          return Arrays.asList(partitions);
      }
  
      protected static TopicPartition tp(String topic, int partition) {
          return new TopicPartition(topic, partition);
      }
  
      protected Optional<String> consumerRackId(int consumerIndex) {
          int numRacks = numBrokerRacks > 0 ? numBrokerRacks : AbstractPartitionAssignorTest.ALL_RACKS.length;
          return Optional.ofNullable(hasConsumerRack ? AbstractPartitionAssignorTest.ALL_RACKS[consumerIndex % numRacks] : null);
      }
  
      protected Subscription subscription(List<String> topics, int consumerIndex) {
          return new Subscription(topics, null, Collections.emptyList(), DEFAULT_GENERATION, consumerRackId(consumerIndex));
      }
  
      protected static boolean isFullyBalanced(Map<String, List<TopicPartition>> assignment) {
          int min = Integer.MAX_VALUE;
          int max = Integer.MIN_VALUE;
          for (List<TopicPartition> topicPartitions: assignment.values()) {
              int size = topicPartitions.size();
              if (size < min)
                  min = size;
              if (size > max)
                  max = size;
          }
          return max - min <= 1;
      }
  
      protected static List<String> getRandomSublist(ArrayList<String> list) {
          List<String> selectedItems = new ArrayList<>(list);
          int len = list.size();
          Random random = new Random();
          int howManyToRemove = random.nextInt(len);
  
          for (int i = 1; i <= howManyToRemove; ++i)
              selectedItems.remove(random.nextInt(selectedItems.size()));
  
          return selectedItems;
      }
  
      /**
       * Verifies that the given assignment is valid with respect to the given subscriptions
       * Validity requirements:
       * - each consumer is subscribed to topics of all partitions assigned to it, and
       * - each partition is assigned to no more than one consumer
       * Balance requirements:
       * - the assignment is fully balanced (the numbers of topic partitions assigned to consumers differ by at most one), or
       * - there is no topic partition that can be moved from one consumer to another with 2+ fewer topic partitions
       *
       * @param subscriptions: topic subscriptions of each consumer
       * @param assignments: given assignment for balance check
       * @param partitionsPerTopic: number of partitions per topic
       */
      protected void verifyValidityAndBalance(Map<String, Subscription> subscriptions,
                                              Map<String, List<TopicPartition>> assignments,
                                              Map<String, List<PartitionInfo>> partitionsPerTopic) {
          int size = subscriptions.size();
          assert size == assignments.size();
  
          List<String> consumers = Utils.sorted(assignments.keySet());
  
          for (int i = 0; i < size; ++i) {
              String consumer = consumers.get(i);
              List<TopicPartition> partitions = assignments.get(consumer);
              for (TopicPartition partition: partitions)
                  assertTrue(subscriptions.get(consumer).topics().contains(partition.topic()),
                      "Error: Partition " + partition + "is assigned to c" + i + ", but it is not subscribed to Topic t" +
                      partition.topic() + "\nSubscriptions: " + subscriptions + "\nAssignments: " + assignments);
  
              if (i == size - 1)
                  continue;
  
              for (int j = i + 1; j < size; ++j) {
                  String otherConsumer = consumers.get(j);
                  List<TopicPartition> otherPartitions = assignments.get(otherConsumer);
  
                  Set<TopicPartition> intersection = new HashSet<>(partitions);
                  intersection.retainAll(otherPartitions);
                  assertTrue(intersection.isEmpty(),
                      "Error: Consumers c" + i + " and c" + j + " have common partitions assigned to them: " + intersection +
                      "\nSubscriptions: " + subscriptions + "\nAssignments: " + assignments);
  
                  int len = partitions.size();
                  int otherLen = otherPartitions.size();
  
                  if (Math.abs(len - otherLen) <= 1)
                      continue;
  
                  Map<String, List<Integer>> map = CollectionUtils.groupPartitionsByTopic(partitions);
                  Map<String, List<Integer>> otherMap = CollectionUtils.groupPartitionsByTopic(otherPartitions);
  
                  int moreLoaded = len > otherLen ? i : j;
                  int lessLoaded = len > otherLen ? j : i;
  
                  // If there's any overlap in the subscribed topics, we should have been able to balance partitions
                  for (String topic: map.keySet()) {
                      assertFalse(otherMap.containsKey(topic),
                          "Error: Some partitions can be moved from c" + moreLoaded + " to c" + lessLoaded + " to achieve a better balance" +
                          "\nc" + i + " has " + len + " partitions, and c" + j + " has " + otherLen + " partitions." +
                          "\nSubscriptions: " + subscriptions +
                          "\nAssignments: " + assignments);
                  }
              }
          }
      }
  
  
      protected AbstractStickyAssignor.MemberData memberData(Subscription subscription) {
          return assignor.memberData(subscription);
      }
  
      protected List<PartitionInfo> partitionInfos(String topic, int numberOfPartitions) {
          int nextIndex = nextPartitionIndex;
          nextPartitionIndex += 1;
          return AbstractPartitionAssignorTest.partitionInfos(topic, numberOfPartitions,
                  replicationFactor, numBrokerRacks, nextIndex);
      }
  
      protected void initializeRacks(RackConfig rackConfig) {
          this.replicationFactor = 3;
          this.numBrokerRacks = rackConfig != RackConfig.NO_BROKER_RACK ? 3 : 0;
          this.hasConsumerRack = rackConfig != RackConfig.NO_CONSUMER_RACK;
          AbstractPartitionAssignorTest.preferRackAwareLogic(assignor, true);
      }
  }
  