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;
  
  import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.Subscription;
  import org.apache.kafka.clients.consumer.internals.AbstractPartitionAssignor;
  import org.apache.kafka.clients.consumer.internals.AbstractPartitionAssignor.MemberInfo;
  import org.apache.kafka.common.TopicPartition;
  
  import org.junit.jupiter.api.Test;
  
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  import java.util.Optional;
  import java.util.stream.Collectors;
  
  import static org.junit.jupiter.api.Assertions.assertEquals;
  import static org.junit.jupiter.api.Assertions.assertTrue;
  
  public class RoundRobinAssignorTest {
  
      private final RoundRobinAssignor assignor = new RoundRobinAssignor();
      private final String topic = "topic";
      private final String consumerId = "consumer";
      private final String topic1 = "topic1";
      private final String topic2 = "topic2";
  
      @Test
      public void testOneConsumerNoTopic() {
          Map<String, Integer> partitionsPerTopic = new HashMap<>();
  
          Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic,
                  Collections.singletonMap(consumerId, new Subscription(Collections.emptyList())));
          assertEquals(Collections.singleton(consumerId), assignment.keySet());
          assertTrue(assignment.get(consumerId).isEmpty());
      }
  
      @Test
      public void testOneConsumerNonexistentTopic() {
          Map<String, Integer> partitionsPerTopic = new HashMap<>();
          Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic,
                  Collections.singletonMap(consumerId, new Subscription(topics(topic))));
  
          assertEquals(Collections.singleton(consumerId), assignment.keySet());
          assertTrue(assignment.get(consumerId).isEmpty());
      }
  
      @Test
      public void testOneConsumerOneTopic() {
          Map<String, Integer> partitionsPerTopic = new HashMap<>();
          partitionsPerTopic.put(topic, 3);
  
          Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic,
                  Collections.singletonMap(consumerId, new Subscription(topics(topic))));
          assertEquals(partitions(tp(topic, 0), tp(topic, 1), tp(topic, 2)), assignment.get(consumerId));
      }
  
      @Test
      public void testOnlyAssignsPartitionsFromSubscribedTopics() {
          String otherTopic = "other";
  
          Map<String, Integer> partitionsPerTopic = new HashMap<>();
          partitionsPerTopic.put(topic, 3);
          partitionsPerTopic.put(otherTopic, 3);
  
          Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic,
                  Collections.singletonMap(consumerId, new Subscription(topics(topic))));
          assertEquals(partitions(tp(topic, 0), tp(topic, 1), tp(topic, 2)), assignment.get(consumerId));
      }
  
      @Test
      public void testOneConsumerMultipleTopics() {
          Map<String, Integer> partitionsPerTopic = setupPartitionsPerTopicWithTwoTopics(1, 2);
  
          Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic,
                  Collections.singletonMap(consumerId, new Subscription(topics(topic1, topic2))));
          assertEquals(partitions(tp(topic1, 0), tp(topic2, 0), tp(topic2, 1)), assignment.get(consumerId));
      }
  
      @Test
      public void testTwoConsumersOneTopicOnePartition() {
          String consumer1 = "consumer1";
          String consumer2 = "consumer2";
  
          Map<String, Integer> partitionsPerTopic = new HashMap<>();
          partitionsPerTopic.put(topic, 1);
  
          Map<String, Subscription> consumers = new HashMap<>();
          consumers.put(consumer1, new Subscription(topics(topic)));
          consumers.put(consumer2, new Subscription(topics(topic)));
  
          Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, consumers);
          assertEquals(partitions(tp(topic, 0)), assignment.get(consumer1));
          assertEquals(Collections.emptyList(), assignment.get(consumer2));
      }
  
      @Test
      public void testTwoConsumersOneTopicTwoPartitions() {
          String consumer1 = "consumer1";
          String consumer2 = "consumer2";
  
          Map<String, Integer> partitionsPerTopic = new HashMap<>();
          partitionsPerTopic.put(topic, 2);
  
          Map<String, Subscription> consumers = new HashMap<>();
          consumers.put(consumer1, new Subscription(topics(topic)));
          consumers.put(consumer2, new Subscription(topics(topic)));
  
          Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, consumers);
          assertEquals(partitions(tp(topic, 0)), assignment.get(consumer1));
          assertEquals(partitions(tp(topic, 1)), assignment.get(consumer2));
      }
  
      @Test
      public void testMultipleConsumersMixedTopics() {
          String topic1 = "topic1";
          String topic2 = "topic2";
          String consumer1 = "consumer1";
          String consumer2 = "consumer2";
          String consumer3 = "consumer3";
  
          Map<String, Integer> partitionsPerTopic = setupPartitionsPerTopicWithTwoTopics(3, 2);
  
          Map<String, Subscription> consumers = new HashMap<>();
          consumers.put(consumer1, new Subscription(topics(topic1)));
          consumers.put(consumer2, new Subscription(topics(topic1, topic2)));
          consumers.put(consumer3, new Subscription(topics(topic1)));
  
          Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, consumers);
          assertEquals(partitions(tp(topic1, 0)), assignment.get(consumer1));
          assertEquals(partitions(tp(topic1, 1), tp(topic2, 0), tp(topic2, 1)), assignment.get(consumer2));
          assertEquals(partitions(tp(topic1, 2)), assignment.get(consumer3));
      }
  
      @Test
      public void testTwoDynamicConsumersTwoTopicsSixPartitions() {
          String topic1 = "topic1";
          String topic2 = "topic2";
          String consumer1 = "consumer1";
          String consumer2 = "consumer2";
  
          Map<String, Integer> partitionsPerTopic = setupPartitionsPerTopicWithTwoTopics(3, 3);
  
          Map<String, Subscription> consumers = new HashMap<>();
          consumers.put(consumer1, new Subscription(topics(topic1, topic2)));
          consumers.put(consumer2, new Subscription(topics(topic1, topic2)));
  
          Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, consumers);
          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));
      }
  
      @Test
      public void testTwoStaticConsumersTwoTopicsSixPartitions() {
          // although consumer 2 has a higher rank than 1, the comparison happens on
          // instance id level.
          String topic1 = "topic1";
          String topic2 = "topic2";
          String consumer1 = "consumer-b";
          String instance1 = "instance1";
          String consumer2 = "consumer-a";
          String instance2 = "instance2";
  
          Map<String, Integer> partitionsPerTopic = setupPartitionsPerTopicWithTwoTopics(3, 3);
  
          Map<String, Subscription> consumers = new HashMap<>();
          Subscription consumer1Subscription = new Subscription(topics(topic1, topic2), null);
          consumer1Subscription.setGroupInstanceId(Optional.of(instance1));
          consumers.put(consumer1, consumer1Subscription);
          Subscription consumer2Subscription = new Subscription(topics(topic1, topic2), null);
          consumer2Subscription.setGroupInstanceId(Optional.of(instance2));
          consumers.put(consumer2, consumer2Subscription);
          Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, consumers);
          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));
      }
  
      @Test
      public void testOneStaticConsumerAndOneDynamicConsumerTwoTopicsSixPartitions() {
          // although consumer 2 has a higher rank than 1, consumer 1 will win the comparison
          // because it has instance id while consumer 2 doesn't.
          String consumer1 = "consumer-b";
          String instance1 = "instance1";
          String consumer2 = "consumer-a";
  
          Map<String, Integer> partitionsPerTopic = setupPartitionsPerTopicWithTwoTopics(3, 3);
  
          Map<String, Subscription> consumers = new HashMap<>();
  
          Subscription consumer1Subscription = new Subscription(topics(topic1, topic2), null);
          consumer1Subscription.setGroupInstanceId(Optional.of(instance1));
          consumers.put(consumer1, consumer1Subscription);
          consumers.put(consumer2, new Subscription(topics(topic1, topic2)));
  
          Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, consumers);
          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));
      }
  
      @Test
      public void testStaticMemberRoundRobinAssignmentPersistent() {
          // Have 3 static members instance1, instance2, instance3 to be persistent
          // across generations. Their assignment shall be the same.
          String consumer1 = "consumer1";
          String instance1 = "instance1";
          String consumer2 = "consumer2";
          String instance2 = "instance2";
          String consumer3 = "consumer3";
          String instance3 = "instance3";
  
          List<MemberInfo> staticMemberInfos = new ArrayList<>();
          staticMemberInfos.add(new MemberInfo(consumer1, Optional.of(instance1)));
          staticMemberInfos.add(new MemberInfo(consumer2, Optional.of(instance2)));
          staticMemberInfos.add(new MemberInfo(consumer3, Optional.of(instance3)));
  
          // Consumer 4 is a dynamic member.
          String consumer4 = "consumer4";
  
          Map<String, Integer> partitionsPerTopic = setupPartitionsPerTopicWithTwoTopics(3, 3);
  
          Map<String, Subscription> consumers = new HashMap<>();
          for (MemberInfo m : staticMemberInfos) {
              Subscription subscription = new Subscription(topics(topic1, topic2), null);
              subscription.setGroupInstanceId(m.groupInstanceId);
              consumers.put(m.memberId, subscription);
          }
          consumers.put(consumer4, new Subscription(topics(topic1, topic2)));
  
          Map<String, List<TopicPartition>> expectedAssignment = new HashMap<>();
          expectedAssignment.put(consumer1, partitions(tp(topic1, 0), tp(topic2, 1)));
          expectedAssignment.put(consumer2, partitions(tp(topic1, 1), tp(topic2, 2)));
          expectedAssignment.put(consumer3, partitions(tp(topic1, 2)));
          expectedAssignment.put(consumer4, partitions(tp(topic2, 0)));
  
          Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, consumers);
          assertEquals(expectedAssignment, assignment);
  
          // Replace dynamic member 4 with a new dynamic member 5.
          consumers.remove(consumer4);
          String consumer5 = "consumer5";
          consumers.put(consumer5, new Subscription(topics(topic1, topic2)));
  
          expectedAssignment.remove(consumer4);
          expectedAssignment.put(consumer5, partitions(tp(topic2, 0)));
          assignment = assignor.assign(partitionsPerTopic, consumers);
          assertEquals(expectedAssignment, assignment);
      }
  
      @Test
      public void testStaticMemberRoundRobinAssignmentPersistentAfterMemberIdChanges() {
          String consumer1 = "consumer1";
          String instance1 = "instance1";
          String consumer2 = "consumer2";
          String instance2 = "instance2";
          String consumer3 = "consumer3";
          String instance3 = "instance3";
          Map<String, String> memberIdToInstanceId = new HashMap<>();
          memberIdToInstanceId.put(consumer1, instance1);
          memberIdToInstanceId.put(consumer2, instance2);
          memberIdToInstanceId.put(consumer3, instance3);
  
          Map<String, Integer> partitionsPerTopic = setupPartitionsPerTopicWithTwoTopics(5, 5);
  
          Map<String, List<TopicPartition>> expectedInstanceAssignment = new HashMap<>();
          expectedInstanceAssignment.put(instance1,
                                         partitions(tp(topic1, 0), tp(topic1, 3), tp(topic2, 1), tp(topic2, 4)));
          expectedInstanceAssignment.put(instance2,
                                         partitions(tp(topic1, 1), tp(topic1, 4), tp(topic2, 2)));
          expectedInstanceAssignment.put(instance3,
                                         partitions(tp(topic1, 2), tp(topic2, 0), tp(topic2, 3)));
  
          List<AbstractPartitionAssignor.MemberInfo> staticMemberInfos = new ArrayList<>();
          for (Map.Entry<String, String> entry : memberIdToInstanceId.entrySet()) {
              staticMemberInfos.add(new AbstractPartitionAssignor.MemberInfo(entry.getKey(), Optional.of(entry.getValue())));
          }
          Map<String, Subscription> consumers = new HashMap<>();
          for (MemberInfo m : staticMemberInfos) {
              Subscription subscription = new Subscription(topics(topic1, topic2), null);
              subscription.setGroupInstanceId(m.groupInstanceId);
              consumers.put(m.memberId, subscription);
          }
  
          Map<String, List<TopicPartition>> staticAssignment =
              checkStaticAssignment(assignor, partitionsPerTopic, consumers);
          assertEquals(expectedInstanceAssignment, staticAssignment);
  
          memberIdToInstanceId.clear();
  
          // Now switch the member.id fields for each member info, the assignment should
          // stay the same as last time.
          String consumer4 = "consumer4";
          String consumer5 = "consumer5";
          consumers.put(consumer4, consumers.get(consumer3));
          consumers.remove(consumer3);
          consumers.put(consumer5, consumers.get(consumer2));
          consumers.remove(consumer2);
          Map<String, List<TopicPartition>> newStaticAssignment =
              checkStaticAssignment(assignor, partitionsPerTopic, consumers);
          assertEquals(staticAssignment, newStaticAssignment);
      }
  
      private static List<String> topics(String... topics) {
          return Arrays.asList(topics);
      }
  
      private static List<TopicPartition> partitions(TopicPartition... partitions) {
          return Arrays.asList(partitions);
      }
  
      private static TopicPartition tp(String topic, int partition) {
          return new TopicPartition(topic, partition);
      }
  
      private Map<String, Integer> setupPartitionsPerTopicWithTwoTopics(int numberOfPartitions1, int numberOfPartitions2) {
          Map<String, Integer> partitionsPerTopic = new HashMap<>();
          partitionsPerTopic.put(topic1, numberOfPartitions1);
          partitionsPerTopic.put(topic2, numberOfPartitions2);
          return partitionsPerTopic;
      }
  
      private static Map<String, List<TopicPartition>> checkStaticAssignment(AbstractPartitionAssignor assignor,
                                                                             Map<String, Integer> partitionsPerTopic,
                                                                             Map<String, Subscription> consumers) {
          Map<String, List<TopicPartition>> assignmentByMemberId = assignor.assign(partitionsPerTopic, consumers);
          return consumers.entrySet().stream()
                  .filter(e -> e.getValue().groupInstanceId().isPresent())
                  .collect(Collectors.toMap(e -> e.getValue().groupInstanceId().get(), e -> assignmentByMemberId.get(e.getKey())));
      }
  }
  