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 kafka.server.share;
  
  import kafka.server.ReplicaManager;
  
  import org.apache.kafka.common.TopicIdPartition;
  import org.apache.kafka.common.TopicPartition;
  import org.apache.kafka.common.Uuid;
  import org.apache.kafka.common.compress.Compression;
  import org.apache.kafka.common.errors.FencedLeaderEpochException;
  import org.apache.kafka.common.message.ShareFetchResponseData;
  import org.apache.kafka.common.message.ShareFetchResponseData.AcquiredRecords;
  import org.apache.kafka.common.protocol.Errors;
  import org.apache.kafka.common.record.FileRecords;
  import org.apache.kafka.common.record.MemoryRecords;
  import org.apache.kafka.common.record.MemoryRecordsBuilder;
  import org.apache.kafka.common.record.RecordBatch;
  import org.apache.kafka.common.record.Records;
  import org.apache.kafka.common.record.SimpleRecord;
  import org.apache.kafka.common.requests.FetchRequest;
  import org.apache.kafka.server.share.SharePartitionKey;
  import org.apache.kafka.server.share.fetch.ShareAcquiredRecords;
  import org.apache.kafka.server.share.fetch.ShareFetch;
  import org.apache.kafka.server.share.fetch.ShareFetchPartitionData;
  import org.apache.kafka.server.storage.log.FetchIsolation;
  import org.apache.kafka.server.storage.log.FetchParams;
  import org.apache.kafka.server.storage.log.FetchPartitionData;
  import org.apache.kafka.storage.internals.log.OffsetResultHolder;
  import org.apache.kafka.storage.log.metrics.BrokerTopicStats;
  import org.apache.kafka.test.TestUtils;
  
  import org.junit.jupiter.api.Test;
  import org.junit.jupiter.api.extension.ExtensionContext;
  import org.junit.jupiter.params.ParameterizedTest;
  import org.junit.jupiter.params.provider.Arguments;
  import org.junit.jupiter.params.provider.ArgumentsProvider;
  import org.junit.jupiter.params.provider.ArgumentsSource;
  import org.mockito.Mockito;
  
  import java.io.IOException;
  import java.nio.ByteBuffer;
  import java.util.LinkedHashMap;
  import java.util.List;
  import java.util.Map;
  import java.util.Optional;
  import java.util.OptionalInt;
  import java.util.OptionalLong;
  import java.util.concurrent.CompletableFuture;
  import java.util.function.BiConsumer;
  import java.util.stream.Stream;
  
  import static org.apache.kafka.server.share.fetch.ShareFetchTestUtils.createFileRecords;
  import static org.apache.kafka.server.share.fetch.ShareFetchTestUtils.createShareAcquiredRecords;
  import static org.apache.kafka.server.share.fetch.ShareFetchTestUtils.memoryRecordsBuilder;
  import static org.junit.jupiter.api.Assertions.assertEquals;
  import static org.junit.jupiter.api.Assertions.assertTrue;
  import static org.mockito.ArgumentMatchers.any;
  import static org.mockito.ArgumentMatchers.anyBoolean;
  import static org.mockito.ArgumentMatchers.anyInt;
  import static org.mockito.ArgumentMatchers.anyLong;
  import static org.mockito.ArgumentMatchers.anyString;
  import static org.mockito.Mockito.doReturn;
  import static org.mockito.Mockito.doThrow;
  import static org.mockito.Mockito.mock;
  import static org.mockito.Mockito.times;
  import static org.mockito.Mockito.when;
  
  public class ShareFetchUtilsTest {
  
      private static final FetchParams FETCH_PARAMS = new FetchParams(
          FetchRequest.ORDINARY_CONSUMER_ID, -1, 0, 1, 1024 * 1024, FetchIsolation.HIGH_WATERMARK,
          Optional.empty(), true);
      private static final int BATCH_SIZE = 500;
      private static final BiConsumer<SharePartitionKey, Throwable> EXCEPTION_HANDLER = (key, exception) -> {
          // No-op
      };
      private static final BrokerTopicStats BROKER_TOPIC_STATS = new BrokerTopicStats();
  
      @Test
      public void testProcessFetchResponse() {
          String groupId = "grp";
          String memberId = Uuid.randomUuid().toString();
          TopicIdPartition tp0 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 0));
          TopicIdPartition tp1 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 1));
  
          SharePartition sp0 = mock(SharePartition.class);
          SharePartition sp1 = mock(SharePartition.class);
  
          when(sp0.nextFetchOffset()).thenReturn((long) 3);
          when(sp1.nextFetchOffset()).thenReturn((long) 3);
  
          when(sp0.acquire(anyString(), anyInt(), anyInt(), anyLong(), any(FetchPartitionData.class), any())).thenReturn(
              createShareAcquiredRecords(new ShareFetchResponseData.AcquiredRecords()
                  .setFirstOffset(0).setLastOffset(3).setDeliveryCount((short) 1)));
          when(sp1.acquire(anyString(), anyInt(), anyInt(), anyLong(), any(FetchPartitionData.class), any())).thenReturn(
              createShareAcquiredRecords(new ShareFetchResponseData.AcquiredRecords()
                  .setFirstOffset(100).setLastOffset(103).setDeliveryCount((short) 1)));
  
          LinkedHashMap<TopicIdPartition, SharePartition> sharePartitions = new LinkedHashMap<>();
          sharePartitions.put(tp0, sp0);
          sharePartitions.put(tp1, sp1);
  
          ShareFetch shareFetch = new ShareFetch(FETCH_PARAMS, groupId, memberId,
              new CompletableFuture<>(), List.of(tp0, tp1), BATCH_SIZE, 100, BROKER_TOPIC_STATS);
  
          MemoryRecords records = MemoryRecords.withRecords(Compression.NONE,
              new SimpleRecord("0".getBytes(), "v".getBytes()),
              new SimpleRecord("1".getBytes(), "v".getBytes()),
              new SimpleRecord("2".getBytes(), "v".getBytes()),
              new SimpleRecord(null, "value".getBytes()));
  
          MemoryRecords records1 = MemoryRecords.withRecords(100L, Compression.NONE,
              new SimpleRecord("0".getBytes(), "v".getBytes()),
              new SimpleRecord("1".getBytes(), "v".getBytes()),
              new SimpleRecord("2".getBytes(), "v".getBytes()),
              new SimpleRecord(null, "value".getBytes()));
  
          List<ShareFetchPartitionData> responseData = List.of(
              new ShareFetchPartitionData(tp0, 0, new FetchPartitionData(Errors.NONE, 0L, 0L,
                  records, Optional.empty(), OptionalLong.empty(), Optional.empty(),
                  OptionalInt.empty(), false)),
              new ShareFetchPartitionData(tp1, 0, new FetchPartitionData(Errors.NONE, 0L, 100L,
                  records1, Optional.empty(), OptionalLong.empty(), Optional.empty(),
                  OptionalInt.empty(), false))
          );
          Map<TopicIdPartition, ShareFetchResponseData.PartitionData> resultData =
                  ShareFetchUtils.processFetchResponse(shareFetch, responseData, sharePartitions, mock(ReplicaManager.class), EXCEPTION_HANDLER);
  
          assertEquals(2, resultData.size());
          assertTrue(resultData.containsKey(tp0));
          assertTrue(resultData.containsKey(tp1));
          assertEquals(0, resultData.get(tp0).partitionIndex());
          assertEquals(1, resultData.get(tp1).partitionIndex());
          assertEquals(Errors.NONE.code(), resultData.get(tp0).errorCode());
          assertEquals(Errors.NONE.code(), resultData.get(tp1).errorCode());
          assertEquals(List.of(new ShareFetchResponseData.AcquiredRecords()
                          .setFirstOffset(0).setLastOffset(3).setDeliveryCount((short) 1)),
                  resultData.get(tp0).acquiredRecords());
          assertEquals(List.of(new ShareFetchResponseData.AcquiredRecords()
                          .setFirstOffset(100).setLastOffset(103).setDeliveryCount((short) 1)),
                  resultData.get(tp1).acquiredRecords());
      }
  
      @Test
      public void testProcessFetchResponseWithEmptyRecords() {
          String groupId = "grp";
          String memberId = Uuid.randomUuid().toString();
          TopicIdPartition tp0 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 0));
          TopicIdPartition tp1 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 1));
  
          SharePartition sp0 = mock(SharePartition.class);
          SharePartition sp1 = mock(SharePartition.class);
  
          when(sp0.nextFetchOffset()).thenReturn((long) 3);
          when(sp1.nextFetchOffset()).thenReturn((long) 3);
  
          when(sp0.acquire(anyString(), anyInt(), anyInt(), anyLong(), any(FetchPartitionData.class), any())).thenReturn(ShareAcquiredRecords.empty());
          when(sp1.acquire(anyString(), anyInt(), anyInt(), anyLong(), any(FetchPartitionData.class), any())).thenReturn(ShareAcquiredRecords.empty());
  
          LinkedHashMap<TopicIdPartition, SharePartition> sharePartitions = new LinkedHashMap<>();
          sharePartitions.put(tp0, sp0);
          sharePartitions.put(tp1, sp1);
  
          ShareFetch shareFetch = new ShareFetch(FETCH_PARAMS, groupId, memberId,
              new CompletableFuture<>(), List.of(tp0, tp1), BATCH_SIZE, 100, BROKER_TOPIC_STATS);
  
          List<ShareFetchPartitionData> responseData = List.of(
              new ShareFetchPartitionData(tp0, 0, new FetchPartitionData(Errors.NONE, 0L, 0L,
                  MemoryRecords.EMPTY, Optional.empty(), OptionalLong.empty(), Optional.empty(),
                  OptionalInt.empty(), false)),
              new ShareFetchPartitionData(tp1, 0, new FetchPartitionData(Errors.NONE, 0L, 0L,
                  MemoryRecords.EMPTY, Optional.empty(), OptionalLong.empty(), Optional.empty(),
                  OptionalInt.empty(), false))
          );
          Map<TopicIdPartition, ShareFetchResponseData.PartitionData> resultData =
              ShareFetchUtils.processFetchResponse(shareFetch, responseData, sharePartitions, mock(ReplicaManager.class), EXCEPTION_HANDLER);
  
          assertEquals(2, resultData.size());
          assertTrue(resultData.containsKey(tp0));
          assertTrue(resultData.containsKey(tp1));
          assertEquals(0, resultData.get(tp0).partitionIndex());
          assertEquals(1, resultData.get(tp1).partitionIndex());
          assertEquals(Errors.NONE.code(), resultData.get(tp0).errorCode());
          assertEquals(Errors.NONE.code(), resultData.get(tp1).errorCode());
          assertEquals(List.of(), resultData.get(tp0).acquiredRecords());
          assertEquals(List.of(), resultData.get(tp1).acquiredRecords());
      }
  
      @Test
      public void testProcessFetchResponseWithLsoMovementForTopicPartition() {
          String groupId = "grp";
  
          TopicIdPartition tp0 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 0));
          TopicIdPartition tp1 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 1));
  
          SharePartition sp0 = Mockito.mock(SharePartition.class);
          SharePartition sp1 = Mockito.mock(SharePartition.class);
  
          LinkedHashMap<TopicIdPartition, SharePartition> sharePartitions = new LinkedHashMap<>();
          sharePartitions.put(tp0, sp0);
          sharePartitions.put(tp1, sp1);
  
          ShareFetch shareFetch = new ShareFetch(FETCH_PARAMS, groupId, Uuid.randomUuid().toString(),
              new CompletableFuture<>(), List.of(tp0, tp1), BATCH_SIZE, 100, BROKER_TOPIC_STATS);
  
          ReplicaManager replicaManager = mock(ReplicaManager.class);
  
          // Mock the replicaManager.fetchOffsetForTimestamp method to return a timestamp and offset for the topic partition.
          FileRecords.TimestampAndOffset timestampAndOffset = new FileRecords.TimestampAndOffset(100L, 1L, Optional.empty());
          doReturn(new OffsetResultHolder(Optional.of(timestampAndOffset), Optional.empty())).when(replicaManager).fetchOffsetForTimestamp(any(TopicPartition.class), anyLong(), any(), any(), anyBoolean());
  
          when(sp0.nextFetchOffset()).thenReturn((long) 0, (long) 5);
          when(sp1.nextFetchOffset()).thenReturn((long) 4, (long) 4);
  
          when(sp0.acquire(anyString(), anyInt(), anyInt(), anyLong(), any(FetchPartitionData.class), any())).thenReturn(
              ShareAcquiredRecords.empty(),
              createShareAcquiredRecords(new ShareFetchResponseData.AcquiredRecords()
                  .setFirstOffset(0).setLastOffset(3).setDeliveryCount((short) 1)));
          when(sp1.acquire(anyString(), anyInt(), anyInt(), anyLong(), any(FetchPartitionData.class), any())).thenReturn(
              createShareAcquiredRecords(new ShareFetchResponseData.AcquiredRecords()
                  .setFirstOffset(100).setLastOffset(103).setDeliveryCount((short) 1)),
              ShareAcquiredRecords.empty());
  
          MemoryRecords records1 = MemoryRecords.withRecords(Compression.NONE,
              new SimpleRecord("0".getBytes(), "v".getBytes()),
              new SimpleRecord("1".getBytes(), "v".getBytes()),
              new SimpleRecord("2".getBytes(), "v".getBytes()),
              new SimpleRecord(null, "value".getBytes()));
  
          List<ShareFetchPartitionData> responseData1 = List.of(
              new ShareFetchPartitionData(tp0, 0, new FetchPartitionData(Errors.OFFSET_OUT_OF_RANGE, 0L, 0L,
                  MemoryRecords.EMPTY, Optional.empty(), OptionalLong.empty(), Optional.empty(),
                  OptionalInt.empty(), false)),
              new ShareFetchPartitionData(tp1, 0, new FetchPartitionData(Errors.NONE, 0L, 0L,
                  records1, Optional.empty(), OptionalLong.empty(), Optional.empty(),
                  OptionalInt.empty(), false))
          );
          Map<TopicIdPartition, ShareFetchResponseData.PartitionData> resultData1 =
              ShareFetchUtils.processFetchResponse(shareFetch, responseData1, sharePartitions, replicaManager, EXCEPTION_HANDLER);
  
          assertEquals(2, resultData1.size());
          assertTrue(resultData1.containsKey(tp0));
          assertTrue(resultData1.containsKey(tp1));
          assertEquals(0, resultData1.get(tp0).partitionIndex());
          assertEquals(1, resultData1.get(tp1).partitionIndex());
          assertEquals(Errors.NONE.code(), resultData1.get(tp0).errorCode());
          assertEquals(Errors.NONE.code(), resultData1.get(tp1).errorCode());
  
          // Since we have OFFSET_OUT_OF_RANGE exception for tp1 and no exception for tp2 from SharePartition class,
          // we should have 1 call for updateCacheAndOffsets for tp0 and 0 calls for tp1.
          Mockito.verify(sp0, times(1)).updateCacheAndOffsets(any(Long.class));
          Mockito.verify(sp1, times(0)).updateCacheAndOffsets(any(Long.class));
  
          MemoryRecords records2 = MemoryRecords.withRecords(100L, Compression.NONE,
              new SimpleRecord("0".getBytes(), "v".getBytes()),
              new SimpleRecord("1".getBytes(), "v".getBytes()),
              new SimpleRecord("2".getBytes(), "v".getBytes()),
              new SimpleRecord(null, "value".getBytes()));
  
          List<ShareFetchPartitionData> responseData2 = List.of(
              new ShareFetchPartitionData(tp0, 0, new FetchPartitionData(Errors.NONE, 0L, 0L,
                  records2, Optional.empty(), OptionalLong.empty(), Optional.empty(),
                  OptionalInt.empty(), false)),
              new ShareFetchPartitionData(tp1, 0, new FetchPartitionData(Errors.NONE, 0L, 0L,
                  MemoryRecords.EMPTY, Optional.empty(), OptionalLong.empty(), Optional.empty(),
                  OptionalInt.empty(), false))
          );
          Map<TopicIdPartition, ShareFetchResponseData.PartitionData> resultData2 =
              ShareFetchUtils.processFetchResponse(shareFetch, responseData2, sharePartitions, replicaManager, EXCEPTION_HANDLER);
  
          assertEquals(2, resultData2.size());
          assertTrue(resultData2.containsKey(tp0));
          assertTrue(resultData2.containsKey(tp1));
          assertEquals(0, resultData2.get(tp0).partitionIndex());
          assertEquals(1, resultData2.get(tp1).partitionIndex());
          assertEquals(Errors.NONE.code(), resultData2.get(tp0).errorCode());
          assertEquals(Errors.NONE.code(), resultData2.get(tp1).errorCode());
  
          // Since we don't see any exception for tp1 and tp2 from SharePartition class,
          // the updateCacheAndOffsets calls should remain the same as the previous case.
          Mockito.verify(sp0, times(1)).updateCacheAndOffsets(1L);
          Mockito.verify(sp1, times(0)).updateCacheAndOffsets(any(Long.class));
      }
  
      @Test
      public void testProcessFetchResponseWhenNoRecordsAreAcquired() {
          String groupId = "grp";
  
          TopicIdPartition tp0 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 0));
  
          SharePartition sp0 = Mockito.mock(SharePartition.class);
          LinkedHashMap<TopicIdPartition, SharePartition> sharePartitions = new LinkedHashMap<>();
          sharePartitions.put(tp0, sp0);
  
          ShareFetch shareFetch = new ShareFetch(FETCH_PARAMS, groupId, Uuid.randomUuid().toString(),
              new CompletableFuture<>(), List.of(tp0), BATCH_SIZE, 100, BROKER_TOPIC_STATS);
  
          ReplicaManager replicaManager = mock(ReplicaManager.class);
  
          // Mock the replicaManager.fetchOffsetForTimestamp method to return a timestamp and offset for the topic partition.
          FileRecords.TimestampAndOffset timestampAndOffset = new FileRecords.TimestampAndOffset(100L, 1L, Optional.empty());
          doReturn(new OffsetResultHolder(Optional.of(timestampAndOffset), Optional.empty())).when(replicaManager).fetchOffsetForTimestamp(any(TopicPartition.class), anyLong(), any(), any(), anyBoolean());
          when(sp0.acquire(anyString(), anyInt(), anyInt(), anyLong(), any(FetchPartitionData.class), any())).thenReturn(ShareAcquiredRecords.empty());
  
          MemoryRecords records = MemoryRecords.withRecords(Compression.NONE,
              new SimpleRecord("0".getBytes(), "v".getBytes()),
              new SimpleRecord("1".getBytes(), "v".getBytes()),
              new SimpleRecord("2".getBytes(), "v".getBytes()),
              new SimpleRecord(null, "value".getBytes()));
  
          // When no records are acquired from share partition.
          List<ShareFetchPartitionData> responseData = List.of(
              new ShareFetchPartitionData(tp0, 0, new FetchPartitionData(Errors.NONE, 0L, 0L,
                  records, Optional.empty(), OptionalLong.empty(), Optional.empty(),
                  OptionalInt.empty(), false)));
  
          Map<TopicIdPartition, ShareFetchResponseData.PartitionData> resultData =
              ShareFetchUtils.processFetchResponse(shareFetch, responseData, sharePartitions, replicaManager, EXCEPTION_HANDLER);
  
          assertEquals(1, resultData.size());
          assertTrue(resultData.containsKey(tp0));
          assertEquals(0, resultData.get(tp0).partitionIndex());
          assertEquals(MemoryRecords.EMPTY, resultData.get(tp0).records());
          assertTrue(resultData.get(tp0).acquiredRecords().isEmpty());
          assertEquals(Errors.NONE.code(), resultData.get(tp0).errorCode());
  
          // When fetch partition data has OFFSET_OUT_OF_RANGE error.
          responseData = List.of(
              new ShareFetchPartitionData(tp0, 0, new FetchPartitionData(Errors.OFFSET_OUT_OF_RANGE, 0L, 0L,
                  records, Optional.empty(), OptionalLong.empty(), Optional.empty(),
                  OptionalInt.empty(), false)));
  
          resultData = ShareFetchUtils.processFetchResponse(shareFetch, responseData, sharePartitions, replicaManager, EXCEPTION_HANDLER);
  
          assertEquals(1, resultData.size());
          assertTrue(resultData.containsKey(tp0));
          assertEquals(0, resultData.get(tp0).partitionIndex());
          assertEquals(MemoryRecords.EMPTY, resultData.get(tp0).records());
          assertTrue(resultData.get(tp0).acquiredRecords().isEmpty());
          assertEquals(Errors.NONE.code(), resultData.get(tp0).errorCode());
  
          Mockito.verify(sp0, times(1)).updateCacheAndOffsets(1L);
      }
  
      @Test
      public void testProcessFetchResponseWithMaxFetchRecords() throws IOException {
          String groupId = "grp";
          TopicIdPartition tp0 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 0));
          TopicIdPartition tp1 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 1));
  
          SharePartition sp0 = Mockito.mock(SharePartition.class);
          SharePartition sp1 = Mockito.mock(SharePartition.class);
  
          LinkedHashMap<TopicIdPartition, SharePartition> sharePartitions = new LinkedHashMap<>();
          sharePartitions.put(tp0, sp0);
          sharePartitions.put(tp1, sp1);
  
          when(sp0.nextFetchOffset()).thenReturn((long) 0, (long) 5);
          when(sp1.nextFetchOffset()).thenReturn((long) 4, (long) 4);
  
          Uuid memberId = Uuid.randomUuid();
          // Set max fetch records to 10
          ShareFetch shareFetch = new ShareFetch(FETCH_PARAMS, groupId, memberId.toString(),
              new CompletableFuture<>(), List.of(tp0, tp1), BATCH_SIZE, 10, BROKER_TOPIC_STATS);
  
          LinkedHashMap<Long, Integer> recordsPerOffset = new LinkedHashMap<>();
          recordsPerOffset.put(0L, 1);
          recordsPerOffset.put(1L, 1);
          recordsPerOffset.put(2L, 1);
          recordsPerOffset.put(3L, 1);
          Records records1 = createFileRecords(recordsPerOffset);
  
          recordsPerOffset.clear();
          recordsPerOffset.put(100L, 4);
          Records records2 = createFileRecords(recordsPerOffset);
  
          FetchPartitionData fetchPartitionData1 = new FetchPartitionData(Errors.NONE, 0L, 0L,
              records1, Optional.empty(), OptionalLong.empty(), Optional.empty(),
              OptionalInt.empty(), false);
          FetchPartitionData fetchPartitionData2 = new FetchPartitionData(Errors.NONE, 0L, 0L,
              records2, Optional.empty(), OptionalLong.empty(), Optional.empty(),
              OptionalInt.empty(), false);
  
          when(sp0.acquire(memberId.toString(), BATCH_SIZE, 10, 0, fetchPartitionData1, FetchIsolation.HIGH_WATERMARK)).thenReturn(
              createShareAcquiredRecords(new ShareFetchResponseData.AcquiredRecords()
                  .setFirstOffset(0).setLastOffset(1).setDeliveryCount((short) 1)));
          when(sp1.acquire(memberId.toString(), BATCH_SIZE, 8, 0, fetchPartitionData2, FetchIsolation.HIGH_WATERMARK)).thenReturn(
              createShareAcquiredRecords(new ShareFetchResponseData.AcquiredRecords()
                  .setFirstOffset(100).setLastOffset(103).setDeliveryCount((short) 1)));
  
          // Send the topic partitions in order so can validate if correct mock is called, accounting
          // the offset count for the acquired records from the previous share partition acquire.
          List<ShareFetchPartitionData> responseData = List.of(
              new ShareFetchPartitionData(tp0, 0, fetchPartitionData1),
              new ShareFetchPartitionData(tp1, 0, fetchPartitionData2)
          );
  
          Map<TopicIdPartition, ShareFetchResponseData.PartitionData> resultData =
              ShareFetchUtils.processFetchResponse(shareFetch, responseData, sharePartitions,
                  mock(ReplicaManager.class), EXCEPTION_HANDLER);
  
          assertEquals(2, resultData.size());
          assertTrue(resultData.containsKey(tp0));
          assertTrue(resultData.containsKey(tp1));
          assertEquals(0, resultData.get(tp0).partitionIndex());
          assertEquals(1, resultData.get(tp1).partitionIndex());
          assertEquals(Errors.NONE.code(), resultData.get(tp0).errorCode());
          assertEquals(Errors.NONE.code(), resultData.get(tp1).errorCode());
          assertEquals(1, resultData.get(tp0).acquiredRecords().size());
          assertEquals(0, resultData.get(tp0).acquiredRecords().get(0).firstOffset());
          assertEquals(1, resultData.get(tp0).acquiredRecords().get(0).lastOffset());
          assertEquals(1, resultData.get(tp1).acquiredRecords().size());
          assertEquals(100, resultData.get(tp1).acquiredRecords().get(0).firstOffset());
          assertEquals(103, resultData.get(tp1).acquiredRecords().get(0).lastOffset());
  
          // Validate the slicing for fetched data happened for tp0 records, not for tp1 records.
          assertTrue(records1.sizeInBytes() > resultData.get(tp0).records().sizeInBytes());
          assertEquals(records2.sizeInBytes(), resultData.get(tp1).records().sizeInBytes());
      }
  
      @Test
      @SuppressWarnings("unchecked")
      public void testProcessFetchResponseWithOffsetFetchException() {
          SharePartition sp0 = Mockito.mock(SharePartition.class);
          when(sp0.leaderEpoch()).thenReturn(1);
  
          TopicIdPartition tp0 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 0));
          LinkedHashMap<TopicIdPartition, SharePartition> sharePartitions = new LinkedHashMap<>();
          sharePartitions.put(tp0, sp0);
  
          ShareFetch shareFetch = mock(ShareFetch.class);
          when(shareFetch.groupId()).thenReturn("grp");
          ReplicaManager replicaManager = mock(ReplicaManager.class);
  
          // Mock the replicaManager.fetchOffsetForTimestamp method to throw exception.
          Throwable exception = new FencedLeaderEpochException("Fenced exception");
          doThrow(exception).when(replicaManager).fetchOffsetForTimestamp(any(TopicPartition.class), anyLong(), any(), any(), anyBoolean());
          when(sp0.acquire(anyString(), anyInt(), anyInt(), anyLong(), any(FetchPartitionData.class), any())).thenReturn(ShareAcquiredRecords.empty());
  
          // When no records are acquired from share partition.
          List<ShareFetchPartitionData> responseData = List.of(
              new ShareFetchPartitionData(tp0, 0, new FetchPartitionData(Errors.OFFSET_OUT_OF_RANGE, 0L, 0L,
                  MemoryRecords.EMPTY, Optional.empty(), OptionalLong.empty(), Optional.empty(),
                  OptionalInt.empty(), false)));
  
          BiConsumer<SharePartitionKey, Throwable> exceptionHandler = mock(BiConsumer.class);
          Map<TopicIdPartition, ShareFetchResponseData.PartitionData> resultData =
              ShareFetchUtils.processFetchResponse(shareFetch, responseData, sharePartitions,
                  replicaManager, exceptionHandler);
  
          assertTrue(resultData.isEmpty());
          Mockito.verify(shareFetch, times(1)).addErroneous(tp0, exception);
          Mockito.verify(exceptionHandler, times(1)).accept(new SharePartitionKey("grp", tp0), exception);
          Mockito.verify(sp0, times(0)).updateCacheAndOffsets(any(Long.class));
      }
  
      @ParameterizedTest(name = "{0}")
      @ArgumentsSource(RecordsArgumentsProvider.class)
      public void testMaybeSliceFetchRecordsSingleBatch(String name, Records records) {
          // Acquire all offsets, should return same records.
          List<AcquiredRecords> acquiredRecords = List.of(new AcquiredRecords().setFirstOffset(5).setLastOffset(14).setDeliveryCount((short) 1));
          Records slicedRecords = ShareFetchUtils.maybeSliceFetchRecords(records, new ShareAcquiredRecords(acquiredRecords, 10));
          assertEquals(records, slicedRecords);
  
          // Acquire offsets out of first offset bound should return the records for the matching batch.
          acquiredRecords = List.of(new AcquiredRecords().setFirstOffset(2).setLastOffset(14).setDeliveryCount((short) 1));
          slicedRecords = ShareFetchUtils.maybeSliceFetchRecords(records, new ShareAcquiredRecords(acquiredRecords, 10));
          assertEquals(records, slicedRecords);
  
          // Acquire offsets out of last offset bound should return the records for the matching batch.
          acquiredRecords = List.of(new AcquiredRecords().setFirstOffset(5).setLastOffset(20).setDeliveryCount((short) 1));
          slicedRecords = ShareFetchUtils.maybeSliceFetchRecords(records, new ShareAcquiredRecords(acquiredRecords, 5));
          assertEquals(records, slicedRecords);
  
          // Acquire only subset of batch offsets, starting from the first offset.
          acquiredRecords = List.of(new AcquiredRecords().setFirstOffset(5).setLastOffset(8).setDeliveryCount((short) 1));
          slicedRecords = ShareFetchUtils.maybeSliceFetchRecords(records, new ShareAcquiredRecords(acquiredRecords, 1));
          assertEquals(records, slicedRecords);
  
          // Acquire only subset of batch offsets, ending at the last offset.
          acquiredRecords = List.of(new AcquiredRecords().setFirstOffset(8).setLastOffset(14).setDeliveryCount((short) 1));
          slicedRecords = ShareFetchUtils.maybeSliceFetchRecords(records, new ShareAcquiredRecords(acquiredRecords, 1));
          assertEquals(records, slicedRecords);
  
          // Acquire only subset of batch offsets, within the batch.
          acquiredRecords = List.of(new AcquiredRecords().setFirstOffset(8).setLastOffset(10).setDeliveryCount((short) 1));
          slicedRecords = ShareFetchUtils.maybeSliceFetchRecords(records, new ShareAcquiredRecords(acquiredRecords, 1));
          assertEquals(records, slicedRecords);
      }
  
      @ParameterizedTest(name = "{0}")
      @ArgumentsSource(MultipleBatchesRecordsArgumentsProvider.class)
      public void testMaybeSliceFetchRecordsMultipleBatches(String name, Records records) {
          // Acquire all offsets, should return same records.
          List<AcquiredRecords> acquiredRecords = List.of(new AcquiredRecords().setFirstOffset(0).setLastOffset(10).setDeliveryCount((short) 1));
          Records slicedRecords = ShareFetchUtils.maybeSliceFetchRecords(records, new ShareAcquiredRecords(acquiredRecords, 11));
          assertEquals(records, slicedRecords);
  
          // Acquire offsets from all batches, but only first record from last batch. Should return
          // all batches.
          acquiredRecords = List.of(new AcquiredRecords().setFirstOffset(0).setLastOffset(7).setDeliveryCount((short) 1));
          slicedRecords = ShareFetchUtils.maybeSliceFetchRecords(records, new ShareAcquiredRecords(acquiredRecords, 5));
          assertEquals(records, slicedRecords);
  
          // Acquire only first batch offsets, should return only first batch.
          acquiredRecords = List.of(new AcquiredRecords().setFirstOffset(0).setLastOffset(2).setDeliveryCount((short) 1));
          slicedRecords = ShareFetchUtils.maybeSliceFetchRecords(records, new ShareAcquiredRecords(acquiredRecords, 5));
          assertTrue(records.sizeInBytes() > slicedRecords.sizeInBytes());
          List<RecordBatch> recordBatches = TestUtils.toList(slicedRecords.batches());
          assertEquals(1, recordBatches.size());
          assertEquals(0, recordBatches.get(0).baseOffset());
          assertEquals(2, recordBatches.get(0).lastOffset());
  
          // Acquire only second batch offsets, should return only second batch.
          acquiredRecords = List.of(new AcquiredRecords().setFirstOffset(3).setLastOffset(4).setDeliveryCount((short) 1));
          slicedRecords = ShareFetchUtils.maybeSliceFetchRecords(records, new ShareAcquiredRecords(acquiredRecords, 5));
          assertTrue(records.sizeInBytes() > slicedRecords.sizeInBytes());
          recordBatches = TestUtils.toList(slicedRecords.batches());
          assertEquals(1, recordBatches.size());
          assertEquals(3, recordBatches.get(0).baseOffset());
          assertEquals(4, recordBatches.get(0).lastOffset());
  
          // Acquire only last batch offsets, should return only last batch.
          acquiredRecords = List.of(new AcquiredRecords().setFirstOffset(7).setLastOffset(10).setDeliveryCount((short) 1));
          slicedRecords = ShareFetchUtils.maybeSliceFetchRecords(records, new ShareAcquiredRecords(acquiredRecords, 1));
          assertTrue(records.sizeInBytes() > slicedRecords.sizeInBytes());
          recordBatches = TestUtils.toList(slicedRecords.batches());
          assertEquals(1, recordBatches.size());
          assertEquals(7, recordBatches.get(0).baseOffset());
          assertEquals(10, recordBatches.get(0).lastOffset());
  
          // Acquire only subset of first batch offsets, should return only first batch.
          acquiredRecords = List.of(new AcquiredRecords().setFirstOffset(1).setLastOffset(1).setDeliveryCount((short) 1));
          slicedRecords = ShareFetchUtils.maybeSliceFetchRecords(records, new ShareAcquiredRecords(acquiredRecords, 1));
          assertTrue(records.sizeInBytes() > slicedRecords.sizeInBytes());
          recordBatches = TestUtils.toList(slicedRecords.batches());
          assertEquals(1, recordBatches.size());
          assertEquals(0, recordBatches.get(0).baseOffset());
          assertEquals(2, recordBatches.get(0).lastOffset());
  
          // Acquire only subset of second batch offsets, should return only second batch.
          acquiredRecords = List.of(new AcquiredRecords().setFirstOffset(4).setLastOffset(4).setDeliveryCount((short) 1));
          slicedRecords = ShareFetchUtils.maybeSliceFetchRecords(records, new ShareAcquiredRecords(acquiredRecords, 1));
          assertTrue(records.sizeInBytes() > slicedRecords.sizeInBytes());
          recordBatches = TestUtils.toList(slicedRecords.batches());
          assertEquals(1, recordBatches.size());
          assertEquals(3, recordBatches.get(0).baseOffset());
          assertEquals(4, recordBatches.get(0).lastOffset());
  
          // Acquire only subset of last batch offsets, should return only last batch.
          acquiredRecords = List.of(new AcquiredRecords().setFirstOffset(8).setLastOffset(8).setDeliveryCount((short) 1));
          slicedRecords = ShareFetchUtils.maybeSliceFetchRecords(records, new ShareAcquiredRecords(acquiredRecords, 1));
          assertTrue(records.sizeInBytes() > slicedRecords.sizeInBytes());
          recordBatches = TestUtils.toList(slicedRecords.batches());
          assertEquals(1, recordBatches.size());
          assertEquals(7, recordBatches.get(0).baseOffset());
          assertEquals(10, recordBatches.get(0).lastOffset());
  
          // Acquire including gaps between batches, should return 2 batches.
          acquiredRecords = List.of(new AcquiredRecords().setFirstOffset(4).setLastOffset(8).setDeliveryCount((short) 1));
          slicedRecords = ShareFetchUtils.maybeSliceFetchRecords(records, new ShareAcquiredRecords(acquiredRecords, 1));
          assertTrue(records.sizeInBytes() > slicedRecords.sizeInBytes());
          recordBatches = TestUtils.toList(slicedRecords.batches());
          assertEquals(2, recordBatches.size());
          assertEquals(3, recordBatches.get(0).baseOffset());
          assertEquals(4, recordBatches.get(0).lastOffset());
          assertEquals(7, recordBatches.get(1).baseOffset());
          assertEquals(10, recordBatches.get(1).lastOffset());
  
          // Acquire with multiple acquired records, should return matching batches.
          acquiredRecords = List.of(
              new AcquiredRecords().setFirstOffset(0).setLastOffset(2).setDeliveryCount((short) 1),
              new AcquiredRecords().setFirstOffset(3).setLastOffset(4).setDeliveryCount((short) 1));
          slicedRecords = ShareFetchUtils.maybeSliceFetchRecords(records, new ShareAcquiredRecords(acquiredRecords, 1));
          assertTrue(records.sizeInBytes() > slicedRecords.sizeInBytes());
          recordBatches = TestUtils.toList(slicedRecords.batches());
          assertEquals(2, recordBatches.size());
          assertEquals(0, recordBatches.get(0).baseOffset());
          assertEquals(2, recordBatches.get(0).lastOffset());
          assertEquals(3, recordBatches.get(1).baseOffset());
          assertEquals(4, recordBatches.get(1).lastOffset());
  
          // Acquire with multiple acquired records of individual offsets from single batch, should return
          // matching batch.
          acquiredRecords = List.of(
              new AcquiredRecords().setFirstOffset(8).setLastOffset(8).setDeliveryCount((short) 1),
              new AcquiredRecords().setFirstOffset(9).setLastOffset(9).setDeliveryCount((short) 1));
          slicedRecords = ShareFetchUtils.maybeSliceFetchRecords(records, new ShareAcquiredRecords(acquiredRecords, 1));
          assertTrue(records.sizeInBytes() > slicedRecords.sizeInBytes());
          recordBatches = TestUtils.toList(slicedRecords.batches());
          assertEquals(1, recordBatches.size());
          assertEquals(7, recordBatches.get(0).baseOffset());
          assertEquals(10, recordBatches.get(0).lastOffset());
  
          // Acquire with multiple acquired records of individual offsets from multiple batch, should return
          // multiple matching batches.
          acquiredRecords = List.of(
              new AcquiredRecords().setFirstOffset(1).setLastOffset(1).setDeliveryCount((short) 1),
              new AcquiredRecords().setFirstOffset(9).setLastOffset(9).setDeliveryCount((short) 1));
          slicedRecords = ShareFetchUtils.maybeSliceFetchRecords(records, new ShareAcquiredRecords(acquiredRecords, 1));
          assertEquals(records.sizeInBytes(), slicedRecords.sizeInBytes());
      }
  
      @ParameterizedTest(name = "{0}")
      @ArgumentsSource(MultipleBatchesRecordsArgumentsProvider.class)
      public void testMaybeSliceFetchRecordsException(String name, Records records) {
          // Send empty acquired records which should trigger an exception and same file records should
          // be returned. The method doesn't expect empty acquired records.
          Records slicedRecords = ShareFetchUtils.maybeSliceFetchRecords(
              records, new ShareAcquiredRecords(List.of(), 3));
          assertEquals(records, slicedRecords);
      }
  
      private static class RecordsArgumentsProvider implements ArgumentsProvider {
          @Override
          public Stream<? extends Arguments> provideArguments(ExtensionContext context) throws Exception {
              return Stream.of(
                  Arguments.of("FileRecords", createFileRecords(Map.of(5L, 10))),
                  Arguments.of("MemoryRecords", createMemoryRecords(5L, 10))
              );
          }
  
          private MemoryRecords createMemoryRecords(long baseOffset, int numRecords) {
              try (MemoryRecordsBuilder recordsBuilder = memoryRecordsBuilder(numRecords, baseOffset)) {
                  return recordsBuilder.build();
              }
          }
      }
  
      private static class MultipleBatchesRecordsArgumentsProvider implements ArgumentsProvider {
          @Override
          public Stream<? extends Arguments> provideArguments(ExtensionContext context) throws Exception {
              LinkedHashMap<Long, Integer> recordsPerOffset = new LinkedHashMap<>();
              recordsPerOffset.put(0L, 3);
              recordsPerOffset.put(3L, 2);
              recordsPerOffset.put(7L, 4); // Gap of 2 offsets between batches.
              return Stream.of(
                  Arguments.of("FileRecords", createFileRecords(recordsPerOffset)),
                  Arguments.of("MemoryRecords", createMemoryRecords(recordsPerOffset))
              );
          }
  
          private MemoryRecords createMemoryRecords(Map<Long, Integer> recordsPerOffset) {
              ByteBuffer buffer = ByteBuffer.allocate(1024);
              recordsPerOffset.forEach((offset, numOfRecords) -> memoryRecordsBuilder(buffer, numOfRecords, offset).close());
              buffer.flip();
  
              return MemoryRecords.readableRecords(buffer);
          }
      }
  }
  