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.Assignment;
  import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.Subscription;
  import org.apache.kafka.common.TopicPartition;
  import org.apache.kafka.common.message.ConsumerProtocolAssignment;
  import org.apache.kafka.common.message.ConsumerProtocolSubscription;
  import org.apache.kafka.common.protocol.types.ArrayOf;
  import org.apache.kafka.common.protocol.types.Field;
  import org.apache.kafka.common.protocol.types.Schema;
  import org.apache.kafka.common.protocol.types.Struct;
  import org.apache.kafka.common.protocol.types.Type;
  
  import org.junit.jupiter.api.Test;
  import org.junit.jupiter.params.ParameterizedTest;
  import org.junit.jupiter.params.provider.ValueSource;
  
  import java.nio.ByteBuffer;
  import java.util.Arrays;
  import java.util.Collections;
  import java.util.List;
  import java.util.Optional;
  import java.util.Set;
  
  import static org.apache.kafka.clients.consumer.internals.AbstractStickyAssignor.DEFAULT_GENERATION;
  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 class ConsumerProtocolTest {
  
      private final TopicPartition tp1 = new TopicPartition("foo", 1);
      private final TopicPartition tp2 = new TopicPartition("bar", 2);
      private final Optional<String> groupInstanceId = Optional.of("instance.id");
      private final int generationId = 1;
      private final Optional<String> rackId = Optional.of("rack-a");
  
      @Test
      public void serializeDeserializeSubscriptionAllVersions() {
          List<TopicPartition> ownedPartitions = Arrays.asList(
              new TopicPartition("foo", 0),
              new TopicPartition("bar", 0));
          Subscription subscription = new Subscription(Arrays.asList("foo", "bar"),
              ByteBuffer.wrap("hello".getBytes()), ownedPartitions, generationId, rackId);
  
          for (short version = ConsumerProtocolSubscription.LOWEST_SUPPORTED_VERSION; version <= ConsumerProtocolSubscription.HIGHEST_SUPPORTED_VERSION; version++) {
              ByteBuffer buffer = ConsumerProtocol.serializeSubscription(subscription, version);
              Subscription parsedSubscription = ConsumerProtocol.deserializeSubscription(buffer);
  
              assertEquals(Set.copyOf(subscription.topics()), Set.copyOf(parsedSubscription.topics()));
              assertEquals(subscription.userData(), parsedSubscription.userData());
              assertFalse(parsedSubscription.groupInstanceId().isPresent());
  
              if (version >= 1) {
                  assertEquals(Set.copyOf(subscription.ownedPartitions()), Set.copyOf(parsedSubscription.ownedPartitions()));
              } else {
                  assertEquals(Collections.emptyList(), parsedSubscription.ownedPartitions());
              }
  
              if (version >= 2) {
                  assertEquals(generationId, parsedSubscription.generationId().orElse(DEFAULT_GENERATION));
              } else {
                  assertFalse(parsedSubscription.generationId().isPresent());
              }
  
              if (version >= 3) {
                  assertEquals(rackId, parsedSubscription.rackId());
              } else {
                  assertEquals(Optional.empty(), parsedSubscription.rackId());
              }
          }
      }
  
      @Test
      public void serializeDeserializeMetadata() {
          Subscription subscription = new Subscription(Arrays.asList("foo", "bar"), ByteBuffer.wrap(new byte[0]));
          ByteBuffer buffer = ConsumerProtocol.serializeSubscription(subscription);
          Subscription parsedSubscription = ConsumerProtocol.deserializeSubscription(buffer);
          assertEquals(Set.copyOf(subscription.topics()), Set.copyOf(parsedSubscription.topics()));
          assertEquals(0, parsedSubscription.userData().limit());
          assertFalse(parsedSubscription.groupInstanceId().isPresent());
          assertFalse(parsedSubscription.generationId().isPresent());
          assertFalse(parsedSubscription.rackId().isPresent());
      }
  
      @Test
      public void serializeDeserializeMetadataAndGroupInstanceId() {
          Subscription subscription = new Subscription(Arrays.asList("foo", "bar"), ByteBuffer.wrap(new byte[0]));
          ByteBuffer buffer = ConsumerProtocol.serializeSubscription(subscription);
  
          Subscription parsedSubscription = ConsumerProtocol.deserializeSubscription(buffer);
          parsedSubscription.setGroupInstanceId(groupInstanceId);
          assertEquals(Set.copyOf(subscription.topics()), Set.copyOf(parsedSubscription.topics()));
          assertEquals(0, parsedSubscription.userData().limit());
          assertEquals(groupInstanceId, parsedSubscription.groupInstanceId());
          assertFalse(parsedSubscription.generationId().isPresent());
          assertFalse(parsedSubscription.rackId().isPresent());
      }
  
      @Test
      public void serializeDeserializeNullSubscriptionUserData() {
          Subscription subscription = new Subscription(Arrays.asList("foo", "bar"), null);
          ByteBuffer buffer = ConsumerProtocol.serializeSubscription(subscription);
          Subscription parsedSubscription = ConsumerProtocol.deserializeSubscription(buffer);
          assertEquals(Set.copyOf(subscription.topics()), Set.copyOf(parsedSubscription.topics()));
          assertNull(parsedSubscription.userData());
          assertFalse(parsedSubscription.rackId().isPresent());
      }
  
      @Test
      public void serializeSubscriptionShouldOrderTopics() {
          assertEquals(
              ConsumerProtocol.serializeSubscription(
                  new Subscription(Arrays.asList("foo", "bar"), null, Arrays.asList(tp1, tp2))
              ),
              ConsumerProtocol.serializeSubscription(
                  new Subscription(Arrays.asList("bar", "foo"), null, Arrays.asList(tp1, tp2))
              )
          );
      }
  
      @Test
      public void serializeSubscriptionShouldOrderOwnedPartitions() {
          assertEquals(
              ConsumerProtocol.serializeSubscription(
                  new Subscription(Arrays.asList("foo", "bar"), null, Arrays.asList(tp1, tp2))
              ),
              ConsumerProtocol.serializeSubscription(
                  new Subscription(Arrays.asList("foo", "bar"), null, Arrays.asList(tp2, tp1))
              )
          );
      }
  
      @Test
      public void deserializeOldSubscriptionVersion() {
          Subscription subscription = new Subscription(Arrays.asList("foo", "bar"), null);
          ByteBuffer buffer = ConsumerProtocol.serializeSubscription(subscription, (short) 0);
          Subscription parsedSubscription = ConsumerProtocol.deserializeSubscription(buffer);
          assertEquals(Set.copyOf(subscription.topics()), Set.copyOf(parsedSubscription.topics()));
          assertNull(parsedSubscription.userData());
          assertTrue(parsedSubscription.ownedPartitions().isEmpty());
          assertFalse(parsedSubscription.generationId().isPresent());
          assertFalse(parsedSubscription.rackId().isPresent());
      }
  
      @ParameterizedTest
      @ValueSource(booleans = {true, false})
      public void deserializeNewSubscriptionWithOldVersion(boolean hasGenerationIdAndRack) {
          Subscription subscription;
          if (hasGenerationIdAndRack) {
              subscription = new Subscription(Arrays.asList("foo", "bar"), null, Collections.singletonList(tp2), generationId, rackId);
          } else {
              subscription = new Subscription(Arrays.asList("foo", "bar"), null, Collections.singletonList(tp2));
          }
          ByteBuffer buffer = ConsumerProtocol.serializeSubscription(subscription);
          // ignore the version assuming it is the old byte code, as it will blindly deserialize as V0
          ConsumerProtocol.deserializeVersion(buffer);
          Subscription parsedSubscription = ConsumerProtocol.deserializeSubscription(buffer, (short) 0);
          assertEquals(Set.copyOf(subscription.topics()), Set.copyOf(parsedSubscription.topics()));
          assertNull(parsedSubscription.userData());
          assertTrue(parsedSubscription.ownedPartitions().isEmpty());
          assertFalse(parsedSubscription.groupInstanceId().isPresent());
          assertFalse(parsedSubscription.generationId().isPresent());
          assertFalse(parsedSubscription.rackId().isPresent());
      }
  
      @Test
      public void deserializeFutureSubscriptionVersion() {
          ByteBuffer buffer = generateFutureSubscriptionVersionData();
  
          Subscription subscription = ConsumerProtocol.deserializeSubscription(buffer);
          subscription.setGroupInstanceId(groupInstanceId);
          assertEquals(Collections.singleton("topic"), Set.copyOf(subscription.topics()));
          assertEquals(Collections.singleton(tp2), Set.copyOf(subscription.ownedPartitions()));
          assertEquals(groupInstanceId, subscription.groupInstanceId());
          assertEquals(generationId, subscription.generationId().orElse(DEFAULT_GENERATION));
          assertEquals(rackId, subscription.rackId());
      }
  
      @Test
      public void serializeDeserializeAssignmentAllVersions() {
          List<TopicPartition> partitions = Arrays.asList(tp1, tp2);
          Assignment assignment = new Assignment(partitions, ByteBuffer.wrap("hello".getBytes()));
  
          for (short version = ConsumerProtocolAssignment.LOWEST_SUPPORTED_VERSION; version <= ConsumerProtocolAssignment.HIGHEST_SUPPORTED_VERSION; version++) {
              ByteBuffer buffer = ConsumerProtocol.serializeAssignment(assignment, version);
              Assignment parsedAssignment = ConsumerProtocol.deserializeAssignment(buffer);
              assertEquals(Set.copyOf(partitions), Set.copyOf(parsedAssignment.partitions()));
              assertEquals(assignment.userData(), parsedAssignment.userData());
          }
      }
  
      @Test
      public void serializeDeserializeConsumerProtocolAssignmentAllVersions() {
          ConsumerProtocolAssignment assignment = new ConsumerProtocolAssignment()
              .setAssignedPartitions(
                  new ConsumerProtocolAssignment.TopicPartitionCollection(Arrays.asList(
                      new ConsumerProtocolAssignment.TopicPartition()
                          .setTopic(tp1.topic())
                          .setPartitions(Collections.singletonList(tp1.partition())),
                      new ConsumerProtocolAssignment.TopicPartition()
                          .setTopic(tp2.topic())
                          .setPartitions(Collections.singletonList(tp2.partition()))
                  ).iterator())
              )
              .setUserData(ByteBuffer.wrap("hello".getBytes()));
  
          for (short version = ConsumerProtocolAssignment.LOWEST_SUPPORTED_VERSION; version <= ConsumerProtocolAssignment.HIGHEST_SUPPORTED_VERSION; version++) {
              ByteBuffer buffer = ConsumerProtocol.serializeAssignment(assignment, version);
              ConsumerProtocolAssignment parsedAssignment = ConsumerProtocol.deserializeConsumerProtocolAssignment(buffer);
              assertEquals(Set.copyOf(assignment.assignedPartitions()), Set.copyOf(parsedAssignment.assignedPartitions()));
              assertEquals(assignment.userData(), parsedAssignment.userData());
          }
      }
  
      @Test
      public void serializeDeserializeAssignment() {
          List<TopicPartition> partitions = Arrays.asList(tp1, tp2);
          ByteBuffer buffer = ConsumerProtocol.serializeAssignment(new Assignment(partitions, ByteBuffer.wrap(new byte[0])));
          Assignment parsedAssignment = ConsumerProtocol.deserializeAssignment(buffer);
          assertEquals(Set.copyOf(partitions), Set.copyOf(parsedAssignment.partitions()));
          assertEquals(0, parsedAssignment.userData().limit());
      }
  
      @Test
      public void deserializeNullAssignmentUserData() {
          List<TopicPartition> partitions = Arrays.asList(tp1, tp2);
          ByteBuffer buffer = ConsumerProtocol.serializeAssignment(new Assignment(partitions, null));
          Assignment parsedAssignment = ConsumerProtocol.deserializeAssignment(buffer);
          assertEquals(Set.copyOf(partitions), Set.copyOf(parsedAssignment.partitions()));
          assertNull(parsedAssignment.userData());
      }
  
      @Test
      public void deserializeFutureAssignmentVersion() {
          // verify that a new version which adds a field is still parseable
          short version = 100;
  
          Schema assignmentSchemaV100 = new Schema(
              new Field("assigned_partitions", new ArrayOf(
                  ConsumerProtocolAssignment.TopicPartition.SCHEMA_0)),
              new Field("user_data", Type.BYTES),
              new Field("foo", Type.STRING));
  
          Struct assignmentV100 = new Struct(assignmentSchemaV100);
          assignmentV100.set("assigned_partitions",
              new Object[]{new Struct(ConsumerProtocolAssignment.TopicPartition.SCHEMA_0)
                  .set("topic", tp1.topic())
                  .set("partitions", new Object[]{tp1.partition()})});
          assignmentV100.set("user_data", ByteBuffer.wrap(new byte[0]));
          assignmentV100.set("foo", "bar");
  
          Struct headerV100 = new Struct(new Schema(new Field("version", Type.INT16)));
          headerV100.set("version", version);
  
          ByteBuffer buffer = ByteBuffer.allocate(assignmentV100.sizeOf() + headerV100.sizeOf());
          headerV100.writeTo(buffer);
          assignmentV100.writeTo(buffer);
  
          buffer.flip();
  
          Assignment assignment = ConsumerProtocol.deserializeAssignment(buffer);
          assertEquals(Set.of(tp1), Set.copyOf(assignment.partitions()));
      }
  
      @Test
      public void serializeDeserializeConsumerProtocolSubscriptionAllVersions() {
          List<TopicPartition> ownedPartitions = Arrays.asList(
              new TopicPartition("foo", 0),
              new TopicPartition("bar", 0));
          Subscription subscription = new Subscription(Arrays.asList("foo", "bar"),
              ByteBuffer.wrap("hello".getBytes()), ownedPartitions, generationId, rackId);
  
          for (short version = ConsumerProtocolSubscription.LOWEST_SUPPORTED_VERSION; version <= ConsumerProtocolSubscription.HIGHEST_SUPPORTED_VERSION; version++) {
              ByteBuffer buffer = ConsumerProtocol.serializeSubscription(subscription, version);
              ConsumerProtocolSubscription parsedSubscription = ConsumerProtocol.deserializeConsumerProtocolSubscription(buffer);
  
              assertEquals(Set.copyOf(subscription.topics()), Set.copyOf(parsedSubscription.topics()));
              assertEquals(subscription.userData(), parsedSubscription.userData());
  
              if (version >= 1) {
                  assertEquals(
                          Set.of(
                              new ConsumerProtocolSubscription.TopicPartition().setTopic("foo").setPartitions(Collections.singletonList(0)),
                              new ConsumerProtocolSubscription.TopicPartition().setTopic("bar").setPartitions(Collections.singletonList(0)
                      )),
                          Set.copyOf(parsedSubscription.ownedPartitions())
                  );
              } else {
                  assertEquals(new ConsumerProtocolSubscription.TopicPartitionCollection(), parsedSubscription.ownedPartitions());
              }
  
              if (version >= 2) {
                  assertEquals(generationId, parsedSubscription.generationId());
              } else {
                  assertEquals(DEFAULT_GENERATION, parsedSubscription.generationId());
              }
  
              if (version >= 3) {
                  assertEquals(rackId.get(), parsedSubscription.rackId());
              } else {
                  assertNull(parsedSubscription.rackId());
              }
          }
      }
  
      private ByteBuffer generateFutureSubscriptionVersionData() {
          // verify that a new version which adds a field is still parseable
          short version = 100;
  
          Schema subscriptionSchemaV100 = new Schema(
              new Field("topics", new ArrayOf(Type.STRING)),
              new Field("user_data", Type.NULLABLE_BYTES),
              new Field("owned_partitions", new ArrayOf(
                  ConsumerProtocolSubscription.TopicPartition.SCHEMA_1)),
              new Field("generation_id", Type.INT32),
              new Field("rack_id", Type.STRING),
              new Field("bar", Type.STRING));
  
          Struct subscriptionV100 = new Struct(subscriptionSchemaV100);
          subscriptionV100.set("topics", new Object[]{"topic"});
          subscriptionV100.set("user_data", ByteBuffer.wrap(new byte[0]));
          subscriptionV100.set("owned_partitions", new Object[]{new Struct(
              ConsumerProtocolSubscription.TopicPartition.SCHEMA_1)
              .set("topic", tp2.topic())
              .set("partitions", new Object[]{tp2.partition()})});
          subscriptionV100.set("generation_id", generationId);
          subscriptionV100.set("rack_id", rackId.orElse(null));
          subscriptionV100.set("bar", "bar");
  
          Struct headerV100 = new Struct(new Schema(new Field("version", Type.INT16)));
          headerV100.set("version", version);
  
          ByteBuffer buffer = ByteBuffer.allocate(subscriptionV100.sizeOf() + headerV100.sizeOf());
          headerV100.writeTo(buffer);
          subscriptionV100.writeTo(buffer);
  
          buffer.flip();
  
          return buffer;
      }
  }
  