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;
  
  import org.apache.kafka.common.TopicIdPartition;
  import org.apache.kafka.common.TopicPartition;
  import org.apache.kafka.common.Uuid;
  import org.apache.kafka.common.message.FetchResponseData;
  import org.apache.kafka.common.protocol.ApiKeys;
  import org.apache.kafka.common.protocol.Errors;
  import org.apache.kafka.common.requests.FetchMetadata;
  import org.apache.kafka.common.requests.FetchRequest;
  import org.apache.kafka.common.requests.FetchResponse;
  import org.apache.kafka.common.utils.LogContext;
  
  import org.junit.jupiter.api.Test;
  import org.junit.jupiter.api.Timeout;
  import org.junit.jupiter.params.ParameterizedTest;
  import org.junit.jupiter.params.provider.Arguments;
  import org.junit.jupiter.params.provider.MethodSource;
  import org.junit.jupiter.params.provider.ValueSource;
  
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.Collections;
  import java.util.Comparator;
  import java.util.HashMap;
  import java.util.Iterator;
  import java.util.LinkedHashMap;
  import java.util.List;
  import java.util.Map;
  import java.util.Optional;
  import java.util.Set;
  import java.util.TreeSet;
  import java.util.stream.Stream;
  
  import static org.apache.kafka.common.requests.FetchMetadata.INITIAL_EPOCH;
  import static org.apache.kafka.common.requests.FetchMetadata.INVALID_SESSION_ID;
  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.assertThrows;
  import static org.junit.jupiter.api.Assertions.assertTrue;
  import static org.junit.jupiter.api.Assertions.fail;
  
  /**
   * A unit test for FetchSessionHandler.
   */
  @Timeout(120)
  public class FetchSessionHandlerTest {
      private static final LogContext LOG_CONTEXT = new LogContext("[FetchSessionHandler]=");
  
      /**
       * Create a set of TopicPartitions.  We use a TreeSet, in order to get a deterministic
       * ordering for test purposes.
       */
      private static Set<TopicPartition> toSet(TopicPartition... arr) {
          TreeSet<TopicPartition> set = new TreeSet<>(Comparator.comparing(TopicPartition::toString));
          set.addAll(Arrays.asList(arr));
          return set;
      }
  
      @Test
      public void testFindMissing() {
          TopicPartition foo0 = new TopicPartition("foo", 0);
          TopicPartition foo1 = new TopicPartition("foo", 1);
          TopicPartition bar0 = new TopicPartition("bar", 0);
          TopicPartition bar1 = new TopicPartition("bar", 1);
          TopicPartition baz0 = new TopicPartition("baz", 0);
          TopicPartition baz1 = new TopicPartition("baz", 1);
          assertEquals(toSet(), FetchSessionHandler.findMissing(toSet(foo0), toSet(foo0)));
          assertEquals(toSet(foo0), FetchSessionHandler.findMissing(toSet(foo0), toSet(foo1)));
          assertEquals(toSet(foo0, foo1),
              FetchSessionHandler.findMissing(toSet(foo0, foo1), toSet(baz0)));
          assertEquals(toSet(bar1, foo0, foo1),
              FetchSessionHandler.findMissing(toSet(foo0, foo1, bar0, bar1),
                  toSet(bar0, baz0, baz1)));
          assertEquals(toSet(),
              FetchSessionHandler.findMissing(toSet(foo0, foo1, bar0, bar1, baz1),
                  toSet(foo0, foo1, bar0, bar1, baz0, baz1)));
      }
  
      private static final class ReqEntry {
          final TopicPartition part;
          final FetchRequest.PartitionData data;
  
          ReqEntry(String topic, Uuid topicId, int partition, long fetchOffset, long logStartOffset, int maxBytes) {
              this.part = new TopicPartition(topic, partition);
              this.data = new FetchRequest.PartitionData(topicId, fetchOffset, logStartOffset, maxBytes, Optional.empty());
          }
      }
  
      private static LinkedHashMap<TopicPartition, FetchRequest.PartitionData> reqMap(ReqEntry... entries) {
          LinkedHashMap<TopicPartition, FetchRequest.PartitionData> map = new LinkedHashMap<>();
          for (ReqEntry entry : entries) {
              map.put(entry.part, entry.data);
          }
          return map;
      }
  
      private static void assertMapEquals(Map<TopicPartition, FetchRequest.PartitionData> expected,
                                          Map<TopicPartition, FetchRequest.PartitionData> actual) {
          Iterator<Map.Entry<TopicPartition, FetchRequest.PartitionData>> expectedIter =
              expected.entrySet().iterator();
          Iterator<Map.Entry<TopicPartition, FetchRequest.PartitionData>> actualIter =
              actual.entrySet().iterator();
          int i = 1;
          while (expectedIter.hasNext()) {
              Map.Entry<TopicPartition, FetchRequest.PartitionData> expectedEntry = expectedIter.next();
              if (!actualIter.hasNext()) {
                  fail("Element " + i + " not found.");
              }
              Map.Entry<TopicPartition, FetchRequest.PartitionData> actuaLEntry = actualIter.next();
              assertEquals(expectedEntry.getKey(), actuaLEntry.getKey(), "Element " + i +
                  " had a different TopicPartition than expected.");
              assertEquals(expectedEntry.getValue(), actuaLEntry.getValue(), "Element " + i +
                  " had different PartitionData than expected.");
              i++;
          }
          if (actualIter.hasNext()) {
              fail("Unexpected element " + i + " found.");
          }
      }
  
      @SafeVarargs
      private static void assertMapsEqual(Map<TopicPartition, FetchRequest.PartitionData> expected,
                                          Map<TopicPartition, FetchRequest.PartitionData>... actuals) {
          for (Map<TopicPartition, FetchRequest.PartitionData> actual : actuals) {
              assertMapEquals(expected, actual);
          }
      }
  
      private static void assertListEquals(List<TopicIdPartition> expected, List<TopicIdPartition> actual) {
          for (TopicIdPartition expectedPart : expected) {
              if (!actual.contains(expectedPart)) {
                  fail("Failed to find expected partition " + expectedPart);
              }
          }
          for (TopicIdPartition actualPart : actual) {
              if (!expected.contains(actualPart)) {
                  fail("Found unexpected partition " + actualPart);
              }
          }
      }
  
      private static final class RespEntry {
          final TopicIdPartition part;
          final FetchResponseData.PartitionData data;
  
          RespEntry(String topic, int partition, Uuid topicId, long highWatermark, long lastStableOffset) {
              this.part = new TopicIdPartition(topicId, new TopicPartition(topic, partition));
  
              this.data = new FetchResponseData.PartitionData()
                  .setPartitionIndex(partition)
                  .setHighWatermark(highWatermark)
                  .setLastStableOffset(lastStableOffset)
                  .setLogStartOffset(0);
          }
  
          RespEntry(String topic, int partition, Uuid topicId, Errors error) {
              this.part = new TopicIdPartition(topicId, new TopicPartition(topic, partition));
  
              this.data = new FetchResponseData.PartitionData()
                      .setPartitionIndex(partition)
                      .setErrorCode(error.code())
                      .setHighWatermark(FetchResponse.INVALID_HIGH_WATERMARK);
          }
      }
  
      private static LinkedHashMap<TopicIdPartition, FetchResponseData.PartitionData> respMap(RespEntry... entries) {
          LinkedHashMap<TopicIdPartition, FetchResponseData.PartitionData> map = new LinkedHashMap<>();
          for (RespEntry entry : entries) {
              map.put(entry.part, entry.data);
          }
          return map;
      }
  
      /**
       * Test the handling of SESSIONLESS responses.
       * Pre-KIP-227 brokers always supply this kind of response.
       */
      @Test
      public void testSessionless() {
          Map<String, Uuid> topicIds = new HashMap<>();
          Map<Uuid, String> topicNames = new HashMap<>();
          // We want to test both on older versions that do not use topic IDs and on newer versions that do.
          List<Short> versions = Arrays.asList((short) 12, ApiKeys.FETCH.latestVersion());
          versions.forEach(version -> {
              FetchSessionHandler handler = new FetchSessionHandler(LOG_CONTEXT, 1);
              FetchSessionHandler.Builder builder = handler.newBuilder();
              addTopicId(topicIds, topicNames, "foo", version);
              Uuid fooId = topicIds.getOrDefault("foo", Uuid.ZERO_UUID);
              builder.add(new TopicPartition("foo", 0),
                      new FetchRequest.PartitionData(fooId, 0, 100, 200, Optional.empty()));
              builder.add(new TopicPartition("foo", 1),
                      new FetchRequest.PartitionData(fooId, 10, 110, 210, Optional.empty()));
              FetchSessionHandler.FetchRequestData data = builder.build();
              assertMapsEqual(reqMap(new ReqEntry("foo", fooId, 0, 0, 100, 200),
                      new ReqEntry("foo", fooId, 1, 10, 110, 210)),
                      data.toSend(), data.sessionPartitions());
              assertEquals(INVALID_SESSION_ID, data.metadata().sessionId());
              assertEquals(INITIAL_EPOCH, data.metadata().epoch());
  
              FetchResponse resp = FetchResponse.of(Errors.NONE, 0, INVALID_SESSION_ID,
                  respMap(new RespEntry("foo", 0, fooId, 0, 0),
                          new RespEntry("foo", 1, fooId, 0, 0)),
                      List.of());
              handler.handleResponse(resp, version);
  
              FetchSessionHandler.Builder builder2 = handler.newBuilder();
              builder2.add(new TopicPartition("foo", 0),
                      new FetchRequest.PartitionData(fooId, 0, 100, 200, Optional.empty()));
              FetchSessionHandler.FetchRequestData data2 = builder2.build();
              assertEquals(INVALID_SESSION_ID, data2.metadata().sessionId());
              assertEquals(INITIAL_EPOCH, data2.metadata().epoch());
              assertMapsEqual(reqMap(new ReqEntry("foo", fooId, 0, 0, 100, 200)),
                      data2.toSend(), data2.sessionPartitions());
          });
      }
  
      /**
       * Test handling an incremental fetch session.
       */
      @Test
      public void testIncrementals() {
          Map<String, Uuid> topicIds = new HashMap<>();
          Map<Uuid, String> topicNames = new HashMap<>();
          // We want to test both on older versions that do not use topic IDs and on newer versions that do.
          List<Short> versions = Arrays.asList((short) 12, ApiKeys.FETCH.latestVersion());
          versions.forEach(version -> {
              FetchSessionHandler handler = new FetchSessionHandler(LOG_CONTEXT, 1);
              FetchSessionHandler.Builder builder = handler.newBuilder();
              addTopicId(topicIds, topicNames, "foo", version);
              Uuid fooId = topicIds.getOrDefault("foo", Uuid.ZERO_UUID);
              TopicPartition foo0 = new TopicPartition("foo", 0);
              TopicPartition foo1 = new TopicPartition("foo", 1);
              builder.add(foo0, new FetchRequest.PartitionData(fooId, 0, 100, 200, Optional.empty()));
              builder.add(foo1, new FetchRequest.PartitionData(fooId, 10, 110, 210, Optional.empty()));
              FetchSessionHandler.FetchRequestData data = builder.build();
              assertMapsEqual(reqMap(new ReqEntry("foo", fooId, 0, 0, 100, 200),
                      new ReqEntry("foo", fooId, 1, 10, 110, 210)),
                      data.toSend(), data.sessionPartitions());
              assertEquals(INVALID_SESSION_ID, data.metadata().sessionId());
              assertEquals(INITIAL_EPOCH, data.metadata().epoch());
  
              FetchResponse resp = FetchResponse.of(Errors.NONE, 0, 123,
                  respMap(new RespEntry("foo", 0, fooId, 10, 20),
                          new RespEntry("foo", 1, fooId, 10, 20)),
                      List.of());
              handler.handleResponse(resp, version);
  
              // Test an incremental fetch request which adds one partition and modifies another.
              FetchSessionHandler.Builder builder2 = handler.newBuilder();
              addTopicId(topicIds, topicNames, "bar", version);
              Uuid barId = topicIds.getOrDefault("bar", Uuid.ZERO_UUID);
              TopicPartition bar0 = new TopicPartition("bar", 0);
              builder2.add(foo0, new FetchRequest.PartitionData(fooId, 0, 100, 200, Optional.empty()));
              builder2.add(foo1, new FetchRequest.PartitionData(fooId, 10, 120, 210, Optional.empty()));
              builder2.add(bar0, new FetchRequest.PartitionData(barId, 20, 200, 200, Optional.empty()));
              FetchSessionHandler.FetchRequestData data2 = builder2.build();
              assertFalse(data2.metadata().isFull());
              assertMapEquals(reqMap(new ReqEntry("foo", fooId, 0, 0, 100, 200),
                      new ReqEntry("foo", fooId, 1, 10, 120, 210),
                      new ReqEntry("bar", barId, 0, 20, 200, 200)),
                      data2.sessionPartitions());
              assertMapEquals(reqMap(new ReqEntry("bar", barId, 0, 20, 200, 200),
                      new ReqEntry("foo", fooId, 1, 10, 120, 210)),
                      data2.toSend());
  
              FetchResponse resp2 = FetchResponse.of(Errors.NONE, 0, 123,
                  respMap(new RespEntry("foo", 1, fooId, 20, 20)),
                      List.of());
              handler.handleResponse(resp2, version);
  
              // Skip building a new request.  Test that handling an invalid fetch session epoch response results
              // in a request which closes the session.
              FetchResponse resp3 = FetchResponse.of(Errors.INVALID_FETCH_SESSION_EPOCH, 0, INVALID_SESSION_ID,
                      respMap(),
                      List.of());
              handler.handleResponse(resp3, version);
  
              FetchSessionHandler.Builder builder4 = handler.newBuilder();
              builder4.add(foo0, new FetchRequest.PartitionData(fooId, 0, 100, 200, Optional.empty()));
              builder4.add(foo1, new FetchRequest.PartitionData(fooId, 10, 120, 210, Optional.empty()));
              builder4.add(bar0, new FetchRequest.PartitionData(barId, 20, 200, 200, Optional.empty()));
              FetchSessionHandler.FetchRequestData data4 = builder4.build();
              assertTrue(data4.metadata().isFull());
              assertEquals(data2.metadata().sessionId(), data4.metadata().sessionId());
              assertEquals(INITIAL_EPOCH, data4.metadata().epoch());
              assertMapsEqual(reqMap(new ReqEntry("foo", fooId, 0, 0, 100, 200),
                      new ReqEntry("foo", fooId, 1, 10, 120, 210),
                      new ReqEntry("bar", barId, 0, 20, 200, 200)),
                      data4.sessionPartitions(), data4.toSend());
          });
      }
  
      /**
       * Test that calling FetchSessionHandler#Builder#build twice fails.
       */
      @Test
      public void testDoubleBuild() {
          FetchSessionHandler handler = new FetchSessionHandler(LOG_CONTEXT, 1);
          FetchSessionHandler.Builder builder = handler.newBuilder();
          builder.add(new TopicPartition("foo", 0),
              new FetchRequest.PartitionData(Uuid.randomUuid(), 0, 100, 200, Optional.empty()));
          builder.build();
          assertThrows(Throwable.class, builder::build, "Expected calling build twice to fail.");
      }
  
      @Test
      public void testIncrementalPartitionRemoval() {
          Map<String, Uuid> topicIds = new HashMap<>();
          Map<Uuid, String> topicNames = new HashMap<>();
          // We want to test both on older versions that do not use topic IDs and on newer versions that do.
          List<Short> versions = Arrays.asList((short) 12, ApiKeys.FETCH.latestVersion());
          versions.forEach(version -> {
              FetchSessionHandler handler = new FetchSessionHandler(LOG_CONTEXT, 1);
              FetchSessionHandler.Builder builder = handler.newBuilder();
              addTopicId(topicIds, topicNames, "foo", version);
              addTopicId(topicIds, topicNames, "bar", version);
              Uuid fooId = topicIds.getOrDefault("foo", Uuid.ZERO_UUID);
              Uuid barId = topicIds.getOrDefault("bar", Uuid.ZERO_UUID);
              TopicPartition foo0 = new TopicPartition("foo", 0);
              TopicPartition foo1 = new TopicPartition("foo", 1);
              TopicPartition bar0 = new TopicPartition("bar", 0);
              builder.add(foo0, new FetchRequest.PartitionData(fooId, 0, 100, 200, Optional.empty()));
              builder.add(foo1, new FetchRequest.PartitionData(fooId, 10, 110, 210, Optional.empty()));
              builder.add(bar0, new FetchRequest.PartitionData(barId, 20, 120, 220, Optional.empty()));
              FetchSessionHandler.FetchRequestData data = builder.build();
              assertMapsEqual(reqMap(new ReqEntry("foo", fooId, 0, 0, 100, 200),
                      new ReqEntry("foo", fooId, 1, 10, 110, 210),
                      new ReqEntry("bar", barId, 0, 20, 120, 220)),
                      data.toSend(), data.sessionPartitions());
              assertTrue(data.metadata().isFull());
  
              FetchResponse resp = FetchResponse.of(Errors.NONE, 0, 123,
                  respMap(new RespEntry("foo", 0, fooId, 10, 20),
                          new RespEntry("foo", 1, fooId, 10, 20),
                          new RespEntry("bar", 0, barId, 10, 20)),
                      List.of());
              handler.handleResponse(resp, version);
  
              // Test an incremental fetch request which removes two partitions.
              FetchSessionHandler.Builder builder2 = handler.newBuilder();
              builder2.add(foo1, new FetchRequest.PartitionData(fooId, 10, 110, 210, Optional.empty()));
              FetchSessionHandler.FetchRequestData data2 = builder2.build();
              assertFalse(data2.metadata().isFull());
              assertEquals(123, data2.metadata().sessionId());
              assertEquals(1, data2.metadata().epoch());
              assertMapEquals(reqMap(new ReqEntry("foo", fooId, 1, 10, 110, 210)),
                      data2.sessionPartitions());
              assertMapEquals(reqMap(), data2.toSend());
              ArrayList<TopicIdPartition> expectedToForget2 = new ArrayList<>();
              expectedToForget2.add(new TopicIdPartition(fooId, foo0));
              expectedToForget2.add(new TopicIdPartition(barId, bar0));
              assertListEquals(expectedToForget2, data2.toForget());
  
              // A FETCH_SESSION_ID_NOT_FOUND response triggers us to close the session.
              // The next request is a session establishing FULL request.
              FetchResponse resp2 = FetchResponse.of(Errors.FETCH_SESSION_ID_NOT_FOUND, 0, INVALID_SESSION_ID, 
                      respMap(), 
                      List.of());
              handler.handleResponse(resp2, version);
  
              FetchSessionHandler.Builder builder3 = handler.newBuilder();
              builder3.add(foo0, new FetchRequest.PartitionData(fooId, 0, 100, 200, Optional.empty()));
              FetchSessionHandler.FetchRequestData data3 = builder3.build();
              assertTrue(data3.metadata().isFull());
              assertEquals(INVALID_SESSION_ID, data3.metadata().sessionId());
              assertEquals(INITIAL_EPOCH, data3.metadata().epoch());
              assertMapsEqual(reqMap(new ReqEntry("foo", fooId, 0, 0, 100, 200)),
                      data3.sessionPartitions(), data3.toSend());
          });
      }
  
      @Test
      public void testTopicIdUsageGrantedOnIdUpgrade() {
          // We want to test adding a topic ID to an existing partition and a new partition in the incremental request.
          // 0 is the existing partition and 1 is the new one.
          List<Integer> partitions = Arrays.asList(0, 1);
          partitions.forEach(partition -> {
              String testType = partition == 0 ? "updating a partition" : "adding a new partition";
              FetchSessionHandler handler = new FetchSessionHandler(LOG_CONTEXT, 1);
              FetchSessionHandler.Builder builder = handler.newBuilder();
              builder.add(new TopicPartition("foo", 0),
                      new FetchRequest.PartitionData(Uuid.ZERO_UUID, 0, 100, 200, Optional.empty()));
              FetchSessionHandler.FetchRequestData data = builder.build();
              assertMapsEqual(reqMap(new ReqEntry("foo", Uuid.ZERO_UUID, 0, 0, 100, 200)),
                      data.toSend(), data.sessionPartitions());
              assertTrue(data.metadata().isFull());
              assertFalse(data.canUseTopicIds());
  
              FetchResponse resp = FetchResponse.of(Errors.NONE, 0, 123,
                      respMap(new RespEntry("foo", 0, Uuid.ZERO_UUID, 10, 20)),
                      List.of());
              handler.handleResponse(resp, (short) 12);
  
              // Try to add a topic ID to an already existing topic partition (0) or a new partition (1) in the session.
              Uuid topicId = Uuid.randomUuid();
              FetchSessionHandler.Builder builder2 = handler.newBuilder();
              builder2.add(new TopicPartition("foo", partition),
                      new FetchRequest.PartitionData(topicId, 10, 110, 210, Optional.empty()));
              FetchSessionHandler.FetchRequestData data2 = builder2.build();
              // Should have the same session ID, and next epoch and can only use topic IDs if the partition was updated.
              boolean updated = partition == 0;
              // The receiving broker will handle closing the session.
              assertEquals(123, data2.metadata().sessionId(), "Did not use same session when " + testType);
              assertEquals(1, data2.metadata().epoch(), "Did not have correct epoch when " + testType);
              assertEquals(updated, data2.canUseTopicIds());
          });
      }
  
      @Test
      public void testIdUsageRevokedOnIdDowngrade() {
          // We want to test removing topic ID from an existing partition and adding a new partition without an ID in the incremental request.
          // 0 is the existing partition and 1 is the new one.
          List<Integer> partitions = Arrays.asList(0, 1);
          partitions.forEach(partition -> {
              String testType = partition == 0 ? "updating a partition" : "adding a new partition";
              Uuid fooId = Uuid.randomUuid();
              FetchSessionHandler handler = new FetchSessionHandler(LOG_CONTEXT, 1);
              FetchSessionHandler.Builder builder = handler.newBuilder();
              builder.add(new TopicPartition("foo", 0),
                      new FetchRequest.PartitionData(fooId, 0, 100, 200, Optional.empty()));
              FetchSessionHandler.FetchRequestData data = builder.build();
              assertMapsEqual(reqMap(new ReqEntry("foo", fooId, 0, 0, 100, 200)),
                      data.toSend(), data.sessionPartitions());
              assertTrue(data.metadata().isFull());
              assertTrue(data.canUseTopicIds());
  
              FetchResponse resp = FetchResponse.of(Errors.NONE, 0, 123,
                      respMap(new RespEntry("foo", 0, fooId, 10, 20)),
                      List.of());
              handler.handleResponse(resp, ApiKeys.FETCH.latestVersion());
  
              // Try to remove a topic ID from an existing topic partition (0) or add a new topic partition (1) without an ID.
              FetchSessionHandler.Builder builder2 = handler.newBuilder();
              builder2.add(new TopicPartition("foo", partition),
                      new FetchRequest.PartitionData(Uuid.ZERO_UUID, 10, 110, 210, Optional.empty()));
              FetchSessionHandler.FetchRequestData data2 = builder2.build();
              // Should have the same session ID, and next epoch and can no longer use topic IDs.
              // The receiving broker will handle closing the session.
              assertEquals(123, data2.metadata().sessionId(), "Did not use same session when " + testType);
              assertEquals(1, data2.metadata().epoch(), "Did not have correct epoch when " + testType);
              assertFalse(data2.canUseTopicIds());
          });
      }
  
      private static Stream<Arguments> idUsageCombinations() {
          return Stream.of(
                  Arguments.of(true, true),
                  Arguments.of(true, false),
                  Arguments.of(false, true),
                  Arguments.of(false, false)
          );
      }
  
      @ParameterizedTest
      @MethodSource("idUsageCombinations")
      public void testTopicIdReplaced(boolean startsWithTopicIds, boolean endsWithTopicIds) {
          TopicPartition tp = new TopicPartition("foo", 0);
          FetchSessionHandler handler = new FetchSessionHandler(LOG_CONTEXT, 1);
          FetchSessionHandler.Builder builder = handler.newBuilder();
          Uuid topicId1 = startsWithTopicIds ? Uuid.randomUuid() : Uuid.ZERO_UUID;
          builder.add(tp, new FetchRequest.PartitionData(topicId1, 0, 100, 200, Optional.empty()));
          FetchSessionHandler.FetchRequestData data = builder.build();
          assertMapsEqual(reqMap(new ReqEntry("foo", topicId1, 0, 0, 100, 200)),
                  data.toSend(), data.sessionPartitions());
          assertTrue(data.metadata().isFull());
          assertEquals(startsWithTopicIds, data.canUseTopicIds());
  
          FetchResponse resp = FetchResponse.of(Errors.NONE, 0, 123, respMap(new RespEntry("foo", 0, topicId1, 10, 20)), List.of());
          short version = startsWithTopicIds ? ApiKeys.FETCH.latestVersion() : 12;
          handler.handleResponse(resp, version);
  
          // Try to add a new topic ID.
          FetchSessionHandler.Builder builder2 = handler.newBuilder();
          Uuid topicId2 = endsWithTopicIds ? Uuid.randomUuid() : Uuid.ZERO_UUID;
          // Use the same data besides the topic ID.
          FetchRequest.PartitionData partitionData = new FetchRequest.PartitionData(topicId2, 0, 100, 200, Optional.empty());
          builder2.add(tp, partitionData);
          FetchSessionHandler.FetchRequestData data2 = builder2.build();
  
          if (startsWithTopicIds && endsWithTopicIds) {
              // If we started with an ID, both a only a new ID will count towards replaced.
              // The old topic ID partition should be in toReplace, and the new one should be in toSend.
              assertEquals(Collections.singletonList(new TopicIdPartition(topicId1, tp)), data2.toReplace());
              assertMapsEqual(reqMap(new ReqEntry("foo", topicId2, 0, 0, 100, 200)),
                      data2.toSend(), data2.sessionPartitions());
  
              // sessionTopicNames should contain only the second topic ID.
              assertEquals(Collections.singletonMap(topicId2, tp.topic()), handler.sessionTopicNames());
  
          } else if (startsWithTopicIds || endsWithTopicIds) {
              // If we downgraded to not using topic IDs we will want to send this data.
              // However, we will not mark the partition as one replaced. In this scenario, we should see the session close due to
              // changing request types.
              // We will have the new topic ID in the session partition map
              assertEquals(Collections.emptyList(), data2.toReplace());
              assertMapsEqual(reqMap(new ReqEntry("foo", topicId2, 0, 0, 100, 200)),
                      data2.toSend(), data2.sessionPartitions());
              // The topicNames map will have the new topic ID if it is valid.
              // The old topic ID should be removed as the map will be empty if the request doesn't use topic IDs.
              if (endsWithTopicIds) {
                  assertEquals(Collections.singletonMap(topicId2, tp.topic()), handler.sessionTopicNames());
              } else {
                  assertEquals(Collections.emptyMap(), handler.sessionTopicNames());
              }
  
          } else {
              // Otherwise, we have no partition in toReplace and since the partition and topic ID was not updated, there is no data to send.
              assertEquals(Collections.emptyList(), data2.toReplace());
              assertEquals(Collections.emptyMap(), data2.toSend());
              assertMapsEqual(reqMap(new ReqEntry("foo", topicId2, 0, 0, 100, 200)), data2.sessionPartitions());
              // There is also nothing in the sessionTopicNames map, as there are no topic IDs used.
              assertEquals(Collections.emptyMap(), handler.sessionTopicNames());
          }
  
          // Should have the same session ID, and next epoch and can use topic IDs if it ended with topic IDs.
          assertEquals(123, data2.metadata().sessionId(), "Did not use same session");
          assertEquals(1, data2.metadata().epoch(), "Did not have correct epoch");
          assertEquals(endsWithTopicIds, data2.canUseTopicIds());
      }
  
      @ParameterizedTest
      @ValueSource(booleans = {true, false})
      public void testSessionEpochWhenMixedUsageOfTopicIDs(boolean startsWithTopicIds) {
          Uuid fooId = startsWithTopicIds ? Uuid.randomUuid() : Uuid.ZERO_UUID;
          Uuid barId = startsWithTopicIds ? Uuid.ZERO_UUID : Uuid.randomUuid();
          short responseVersion = startsWithTopicIds ? ApiKeys.FETCH.latestVersion() : 12;
  
          TopicPartition tp0 = new TopicPartition("foo", 0);
          TopicPartition tp1 = new TopicPartition("bar", 1);
  
          FetchSessionHandler handler = new FetchSessionHandler(LOG_CONTEXT, 1);
          FetchSessionHandler.Builder builder = handler.newBuilder();
          builder.add(tp0, new FetchRequest.PartitionData(fooId, 0, 100, 200, Optional.empty()));
          FetchSessionHandler.FetchRequestData data = builder.build();
          assertMapsEqual(reqMap(new ReqEntry("foo", fooId, 0, 0, 100, 200)),
                  data.toSend(), data.sessionPartitions());
          assertTrue(data.metadata().isFull());
          assertEquals(startsWithTopicIds, data.canUseTopicIds());
  
          FetchResponse resp = FetchResponse.of(Errors.NONE, 0, 123,
                  respMap(new RespEntry("foo", 0, fooId, 10, 20)), List.of());
          handler.handleResponse(resp, responseVersion);
  
          // Re-add the first partition. Then add a partition with opposite ID usage.
          FetchSessionHandler.Builder builder2 = handler.newBuilder();
          builder2.add(tp0, new FetchRequest.PartitionData(fooId, 10, 110, 210, Optional.empty()));
          builder2.add(tp1, new FetchRequest.PartitionData(barId, 0, 100, 200, Optional.empty()));
          FetchSessionHandler.FetchRequestData data2 = builder2.build();
          // Should have the same session ID, and the next epoch and can not use topic IDs.
          // The receiving broker will handle closing the session.
          assertEquals(123, data2.metadata().sessionId(), "Did not use same session");
          assertEquals(1, data2.metadata().epoch(), "Did not have final epoch");
          assertFalse(data2.canUseTopicIds());
      }
  
  
      @ParameterizedTest
      @ValueSource(booleans = {true, false})
      public void testIdUsageWithAllForgottenPartitions(boolean useTopicIds) {
          // We want to test when all topics are removed from the session
          TopicPartition foo0 = new TopicPartition("foo", 0);
          Uuid topicId = useTopicIds ? Uuid.randomUuid() : Uuid.ZERO_UUID;
          short responseVersion = useTopicIds ? ApiKeys.FETCH.latestVersion() : 12;
          FetchSessionHandler handler = new FetchSessionHandler(LOG_CONTEXT, 1);
  
          // Add topic foo to the session
          FetchSessionHandler.Builder builder = handler.newBuilder();
          builder.add(foo0, new FetchRequest.PartitionData(topicId, 0, 100, 200, Optional.empty()));
          FetchSessionHandler.FetchRequestData data = builder.build();
          assertMapsEqual(reqMap(new ReqEntry("foo", topicId, 0, 0, 100, 200)),
                  data.toSend(), data.sessionPartitions());
          assertTrue(data.metadata().isFull());
          assertEquals(useTopicIds, data.canUseTopicIds());
  
          FetchResponse resp = FetchResponse.of(Errors.NONE, 0, 123,
                  respMap(new RespEntry("foo", 0, topicId, 10, 20)),
                  List.of());
          handler.handleResponse(resp, responseVersion);
  
          // Remove the topic from the session
          FetchSessionHandler.Builder builder2 = handler.newBuilder();
          FetchSessionHandler.FetchRequestData data2 = builder2.build();
          assertEquals(Collections.singletonList(new TopicIdPartition(topicId, foo0)), data2.toForget());
          // Should have the same session ID, next epoch, and same ID usage.
          assertEquals(123, data2.metadata().sessionId(), "Did not use same session when useTopicIds was " + useTopicIds);
          assertEquals(1, data2.metadata().epoch(), "Did not have correct epoch when useTopicIds was " + useTopicIds);
          assertEquals(useTopicIds, data2.canUseTopicIds());
      }
  
      @Test
      public void testOkToAddNewIdAfterTopicRemovedFromSession() {
          Uuid topicId = Uuid.randomUuid();
          FetchSessionHandler handler = new FetchSessionHandler(LOG_CONTEXT, 1);
          FetchSessionHandler.Builder builder = handler.newBuilder();
          builder.add(new TopicPartition("foo", 0),
                  new FetchRequest.PartitionData(topicId, 0, 100, 200, Optional.empty()));
          FetchSessionHandler.FetchRequestData data = builder.build();
          assertMapsEqual(reqMap(new ReqEntry("foo", topicId, 0, 0, 100, 200)),
                  data.toSend(), data.sessionPartitions());
          assertTrue(data.metadata().isFull());
          assertTrue(data.canUseTopicIds());
  
          FetchResponse resp = FetchResponse.of(Errors.NONE, 0, 123,
                  respMap(new RespEntry("foo", 0, topicId, 10, 20)),
                  List.of());
          handler.handleResponse(resp, ApiKeys.FETCH.latestVersion());
  
          // Remove the partition from the session. Return a session ID as though the session is still open.
          FetchSessionHandler.Builder builder2 = handler.newBuilder();
          FetchSessionHandler.FetchRequestData data2 = builder2.build();
          assertMapsEqual(new LinkedHashMap<>(),
                  data2.toSend(), data2.sessionPartitions());
          FetchResponse resp2 = FetchResponse.of(Errors.NONE, 0, 123,
                  new LinkedHashMap<>(),
                  List.of());
          handler.handleResponse(resp2, ApiKeys.FETCH.latestVersion());
  
          // After the topic is removed, add a recreated topic with a new ID.
          FetchSessionHandler.Builder builder3 = handler.newBuilder();
          builder3.add(new TopicPartition("foo", 0),
                  new FetchRequest.PartitionData(Uuid.randomUuid(), 0, 100, 200, Optional.empty()));
          FetchSessionHandler.FetchRequestData data3 = builder3.build();
          // Should have the same session ID and epoch 2.
          assertEquals(123, data3.metadata().sessionId(), "Did not use same session");
          assertEquals(2, data3.metadata().epoch(), "Did not have the correct session epoch");
          assertTrue(data.canUseTopicIds());
      }
  
      @Test
      public void testVerifyFullFetchResponsePartitions() {
          Map<String, Uuid> topicIds = new HashMap<>();
          Map<Uuid, String> topicNames = new HashMap<>();
          // We want to test both on older versions that do not use topic IDs and on newer versions that do.
          List<Short> versions = Arrays.asList((short) 12, ApiKeys.FETCH.latestVersion());
          versions.forEach(version -> {
              FetchSessionHandler handler = new FetchSessionHandler(LOG_CONTEXT, 1);
              addTopicId(topicIds, topicNames, "foo", version);
              addTopicId(topicIds, topicNames, "bar", version);
              Uuid fooId = topicIds.getOrDefault("foo", Uuid.ZERO_UUID);
              Uuid barId = topicIds.getOrDefault("bar", Uuid.ZERO_UUID);
              TopicPartition foo0 = new TopicPartition("foo", 0);
              TopicPartition foo1 = new TopicPartition("foo", 1);
              TopicPartition bar0 = new TopicPartition("bar", 0);
              FetchResponse resp1 = FetchResponse.of(Errors.NONE, 0, INVALID_SESSION_ID,
                  respMap(new RespEntry("foo", 0, fooId, 10, 20),
                          new RespEntry("foo", 1, fooId, 10, 20),
                          new RespEntry("bar", 0, barId, 10, 20)),
                      List.of());
              String issue = handler.verifyFullFetchResponsePartitions(resp1.responseData(topicNames, version).keySet(),
                      resp1.topicIds(), version);
              assertTrue(issue.contains("extraPartitions="));
              assertFalse(issue.contains("omittedPartitions="));
              FetchSessionHandler.Builder builder = handler.newBuilder();
              builder.add(foo0, new FetchRequest.PartitionData(fooId, 0, 100, 200, Optional.empty()));
              builder.add(foo1, new FetchRequest.PartitionData(fooId, 10, 110, 210, Optional.empty()));
              builder.add(bar0, new FetchRequest.PartitionData(barId, 20, 120, 220, Optional.empty()));
              builder.build();
              FetchResponse resp2 = FetchResponse.of(Errors.NONE, 0, INVALID_SESSION_ID,
                  respMap(new RespEntry("foo", 0, fooId, 10, 20),
                          new RespEntry("foo", 1, fooId, 10, 20),
                          new RespEntry("bar", 0, barId, 10, 20)),
                      List.of());
              String issue2 = handler.verifyFullFetchResponsePartitions(resp2.responseData(topicNames, version).keySet(),
                      resp2.topicIds(), version);
              assertNull(issue2);
              FetchResponse resp3 = FetchResponse.of(Errors.NONE, 0, INVALID_SESSION_ID,
                  respMap(new RespEntry("foo", 0, fooId, 10, 20),
                          new RespEntry("foo", 1, fooId, 10, 20)),
                      List.of());
              String issue3 = handler.verifyFullFetchResponsePartitions(resp3.responseData(topicNames, version).keySet(),
                      resp3.topicIds(), version);
              assertFalse(issue3.contains("extraPartitions="));
              assertTrue(issue3.contains("omittedPartitions="));
          });
      }
  
      @Test
      public void testVerifyFullFetchResponsePartitionsWithTopicIds() {
          Map<String, Uuid> topicIds = new HashMap<>();
          Map<Uuid, String> topicNames = new HashMap<>();
          FetchSessionHandler handler = new FetchSessionHandler(LOG_CONTEXT, 1);
          addTopicId(topicIds, topicNames, "foo", ApiKeys.FETCH.latestVersion());
          addTopicId(topicIds, topicNames, "bar", ApiKeys.FETCH.latestVersion());
          addTopicId(topicIds, topicNames, "extra2", ApiKeys.FETCH.latestVersion());
          FetchResponse resp1 = FetchResponse.of(Errors.NONE, 0, INVALID_SESSION_ID,
              respMap(new RespEntry("foo", 0, topicIds.get("foo"), 10, 20),
                      new RespEntry("extra2", 1, topicIds.get("extra2"), 10, 20),
                      new RespEntry("bar", 0, topicIds.get("bar"), 10, 20)),
                  List.of());
          String issue = handler.verifyFullFetchResponsePartitions(resp1.responseData(topicNames, ApiKeys.FETCH.latestVersion()).keySet(),
                  resp1.topicIds(), ApiKeys.FETCH.latestVersion());
          assertTrue(issue.contains("extraPartitions="));
          assertFalse(issue.contains("omittedPartitions="));
          FetchSessionHandler.Builder builder = handler.newBuilder();
          builder.add(new TopicPartition("foo", 0),
                  new FetchRequest.PartitionData(topicIds.get("foo"), 0, 100, 200, Optional.empty()));
          builder.add(new TopicPartition("bar", 0),
                  new FetchRequest.PartitionData(topicIds.get("bar"), 20, 120, 220, Optional.empty()));
          builder.build();
          FetchResponse resp2 = FetchResponse.of(Errors.NONE, 0, INVALID_SESSION_ID,
              respMap(new RespEntry("foo", 0, topicIds.get("foo"), 10, 20),
                      new RespEntry("extra2", 1, topicIds.get("extra2"), 10, 20),
                      new RespEntry("bar", 0, topicIds.get("bar"), 10, 20)),
                  List.of());
          String issue2 = handler.verifyFullFetchResponsePartitions(resp2.responseData(topicNames, ApiKeys.FETCH.latestVersion()).keySet(),
                  resp2.topicIds(), ApiKeys.FETCH.latestVersion());
          assertTrue(issue2.contains("extraPartitions="));
          assertFalse(issue2.contains("omittedPartitions="));
          FetchResponse resp3 = FetchResponse.of(Errors.NONE, 0, INVALID_SESSION_ID,
              respMap(new RespEntry("foo", 0, topicIds.get("foo"), 10, 20),
                      new RespEntry("bar", 0, topicIds.get("bar"), 10, 20)),
                  List.of());
          String issue3 = handler.verifyFullFetchResponsePartitions(resp3.responseData(topicNames, ApiKeys.FETCH.latestVersion()).keySet(),
                  resp3.topicIds(), ApiKeys.FETCH.latestVersion());
          assertNull(issue3);
      }
  
      @Test
      public void testTopLevelErrorResetsMetadata() {
          Map<String, Uuid> topicIds = new HashMap<>();
          Map<Uuid, String> topicNames = new HashMap<>();
          FetchSessionHandler handler = new FetchSessionHandler(LOG_CONTEXT, 1);
          FetchSessionHandler.Builder builder = handler.newBuilder();
          addTopicId(topicIds, topicNames, "foo", ApiKeys.FETCH.latestVersion());
          Uuid fooId = topicIds.getOrDefault("foo", Uuid.ZERO_UUID);
          builder.add(new TopicPartition("foo", 0),
                  new FetchRequest.PartitionData(fooId, 0, 100, 200, Optional.empty()));
          builder.add(new TopicPartition("foo", 1),
                  new FetchRequest.PartitionData(fooId, 10, 110, 210, Optional.empty()));
          FetchSessionHandler.FetchRequestData data = builder.build();
          assertEquals(INVALID_SESSION_ID, data.metadata().sessionId());
          assertEquals(INITIAL_EPOCH, data.metadata().epoch());
  
          FetchResponse resp = FetchResponse.of(Errors.NONE, 0, 123,
              respMap(new RespEntry("foo", 0, topicIds.get("foo"), 10, 20),
                      new RespEntry("foo", 1, topicIds.get("foo"), 10, 20)),
                  List.of());
          handler.handleResponse(resp, ApiKeys.FETCH.latestVersion());
  
          // Test an incremental fetch request which adds an ID unknown to the broker.
          FetchSessionHandler.Builder builder2 = handler.newBuilder();
          addTopicId(topicIds, topicNames, "unknown", ApiKeys.FETCH.latestVersion());
          builder2.add(new TopicPartition("unknown", 0),
                  new FetchRequest.PartitionData(topicIds.getOrDefault("unknown", Uuid.ZERO_UUID), 0, 100, 200, Optional.empty()));
          FetchSessionHandler.FetchRequestData data2 = builder2.build();
          assertFalse(data2.metadata().isFull());
          assertEquals(123, data2.metadata().sessionId());
          assertEquals(FetchMetadata.nextEpoch(INITIAL_EPOCH), data2.metadata().epoch());
  
          // Return and handle a response with a top level error
          FetchResponse resp2 = FetchResponse.of(Errors.UNKNOWN_TOPIC_ID, 0, 123,
              respMap(new RespEntry("unknown", 0, Uuid.randomUuid(), Errors.UNKNOWN_TOPIC_ID)),
                  List.of());
          assertFalse(handler.handleResponse(resp2, ApiKeys.FETCH.latestVersion()));
  
          // Ensure we start with a new epoch. This will close the session in the next request.
          FetchSessionHandler.Builder builder3 = handler.newBuilder();
          FetchSessionHandler.FetchRequestData data3 = builder3.build();
          assertEquals(123, data3.metadata().sessionId());
          assertEquals(INITIAL_EPOCH, data3.metadata().epoch());
      }
  
      private void addTopicId(Map<String, Uuid> topicIds, Map<Uuid, String> topicNames, String name, short version) {
          if (version >= 13) {
              Uuid id = Uuid.randomUuid();
              topicIds.put(name, id);
              topicNames.put(id, name);
          }
      }
  }
  