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.cluster.Partition;
  import kafka.server.ReplicaManager;
  import kafka.server.ReplicaQuota;
  import kafka.server.share.SharePartitionManager.SharePartitionListener;
  
  import org.apache.kafka.clients.consumer.AcknowledgeType;
  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.CoordinatorNotAvailableException;
  import org.apache.kafka.common.errors.FencedStateEpochException;
  import org.apache.kafka.common.errors.InvalidRecordStateException;
  import org.apache.kafka.common.errors.InvalidRequestException;
  import org.apache.kafka.common.errors.InvalidShareSessionEpochException;
  import org.apache.kafka.common.errors.KafkaStorageException;
  import org.apache.kafka.common.errors.LeaderNotAvailableException;
  import org.apache.kafka.common.errors.NotLeaderOrFollowerException;
  import org.apache.kafka.common.errors.ShareSessionLimitReachedException;
  import org.apache.kafka.common.errors.ShareSessionNotFoundException;
  import org.apache.kafka.common.message.ShareAcknowledgeResponseData;
  import org.apache.kafka.common.message.ShareFetchResponseData;
  import org.apache.kafka.common.message.ShareFetchResponseData.PartitionData;
  import org.apache.kafka.common.protocol.ApiKeys;
  import org.apache.kafka.common.protocol.Errors;
  import org.apache.kafka.common.protocol.ObjectSerializationCache;
  import org.apache.kafka.common.record.FileRecords;
  import org.apache.kafka.common.record.MemoryRecords;
  import org.apache.kafka.common.record.SimpleRecord;
  import org.apache.kafka.common.requests.FetchRequest;
  import org.apache.kafka.common.requests.ShareFetchResponse;
  import org.apache.kafka.common.requests.ShareRequestMetadata;
  import org.apache.kafka.common.utils.ImplicitLinkedHashCollection;
  import org.apache.kafka.common.utils.Time;
  import org.apache.kafka.coordinator.group.GroupConfigManager;
  import org.apache.kafka.server.LogReadResult;
  import org.apache.kafka.server.common.ShareVersion;
  import org.apache.kafka.server.purgatory.DelayedOperationKey;
  import org.apache.kafka.server.purgatory.DelayedOperationPurgatory;
  import org.apache.kafka.server.share.CachedSharePartition;
  import org.apache.kafka.server.share.ErroneousAndValidPartitionData;
  import org.apache.kafka.server.share.SharePartitionKey;
  import org.apache.kafka.server.share.acknowledge.ShareAcknowledgementBatch;
  import org.apache.kafka.server.share.context.FinalContext;
  import org.apache.kafka.server.share.context.ShareFetchContext;
  import org.apache.kafka.server.share.context.ShareSessionContext;
  import org.apache.kafka.server.share.fetch.DelayedShareFetchGroupKey;
  import org.apache.kafka.server.share.fetch.DelayedShareFetchKey;
  import org.apache.kafka.server.share.fetch.PartitionMaxBytesStrategy;
  import org.apache.kafka.server.share.fetch.ShareAcquiredRecords;
  import org.apache.kafka.server.share.fetch.ShareFetch;
  import org.apache.kafka.server.share.metrics.ShareGroupMetrics;
  import org.apache.kafka.server.share.persister.NoOpStatePersister;
  import org.apache.kafka.server.share.persister.Persister;
  import org.apache.kafka.server.share.session.ShareSession;
  import org.apache.kafka.server.share.session.ShareSessionCache;
  import org.apache.kafka.server.share.session.ShareSessionKey;
  import org.apache.kafka.server.storage.log.FetchIsolation;
  import org.apache.kafka.server.storage.log.FetchParams;
  import org.apache.kafka.server.util.FutureUtils;
  import org.apache.kafka.server.util.MockTime;
  import org.apache.kafka.server.util.timer.MockTimer;
  import org.apache.kafka.server.util.timer.SystemTimer;
  import org.apache.kafka.server.util.timer.SystemTimerReaper;
  import org.apache.kafka.server.util.timer.Timer;
  import org.apache.kafka.server.util.timer.TimerTask;
  import org.apache.kafka.storage.internals.log.FetchDataInfo;
  import org.apache.kafka.storage.internals.log.LogOffsetMetadata;
  import org.apache.kafka.storage.internals.log.OffsetResultHolder;
  import org.apache.kafka.storage.log.metrics.BrokerTopicMetrics;
  import org.apache.kafka.storage.log.metrics.BrokerTopicStats;
  import org.apache.kafka.test.TestUtils;
  
  import org.junit.jupiter.api.AfterEach;
  import org.junit.jupiter.api.BeforeEach;
  import org.junit.jupiter.api.Test;
  import org.junit.jupiter.api.Timeout;
  import org.mockito.ArgumentCaptor;
  import org.mockito.ArgumentMatchers;
  import org.mockito.Mockito;
  
  import java.util.ArrayList;
  import java.util.HashMap;
  import java.util.HashSet;
  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.Set;
  import java.util.concurrent.CompletableFuture;
  import java.util.concurrent.ExecutionException;
  import java.util.function.Consumer;
  
  import scala.Tuple2;
  import scala.collection.Seq;
  import scala.jdk.javaapi.CollectionConverters;
  
  import static kafka.server.share.DelayedShareFetchTest.mockTopicIdPartitionToReturnDataEqualToMinBytes;
  import static org.apache.kafka.server.share.fetch.ShareFetchTestUtils.validateRotatedListEquals;
  import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
  import static org.junit.jupiter.api.Assertions.assertEquals;
  import static org.junit.jupiter.api.Assertions.assertFalse;
  import static org.junit.jupiter.api.Assertions.assertInstanceOf;
  import static org.junit.jupiter.api.Assertions.assertNotEquals;
  import static org.junit.jupiter.api.Assertions.assertNotNull;
  import static org.junit.jupiter.api.Assertions.assertThrows;
  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.doAnswer;
  import static org.mockito.Mockito.doThrow;
  import static org.mockito.Mockito.mock;
  import static org.mockito.Mockito.spy;
  import static org.mockito.Mockito.times;
  import static org.mockito.Mockito.verify;
  import static org.mockito.Mockito.when;
  
  @Timeout(120)
  @SuppressWarnings({"ClassDataAbstractionCoupling"})
  public class SharePartitionManagerTest {
  
      private static final int DEFAULT_RECORD_LOCK_DURATION_MS = 30000;
      private static final int MAX_DELIVERY_COUNT = 5;
      private static final short MAX_IN_FLIGHT_MESSAGES = 200;
      private static final short MAX_FETCH_RECORDS = 500;
      private static final int DELAYED_SHARE_FETCH_MAX_WAIT_MS = 2000;
      private static final int DELAYED_SHARE_FETCH_TIMEOUT_MS = 3000;
      private static final int BATCH_SIZE = 500;
      private static final FetchParams FETCH_PARAMS = new FetchParams(
          FetchRequest.ORDINARY_CONSUMER_ID, -1, DELAYED_SHARE_FETCH_MAX_WAIT_MS,
          1, 1024 * 1024, FetchIsolation.HIGH_WATERMARK, Optional.empty(), true);
      private static final String TIMER_NAME_PREFIX = "share-partition-manager";
      private static final String CONNECTION_ID = "id-1";
  
      static final int DELAYED_SHARE_FETCH_PURGATORY_PURGE_INTERVAL = 1000;
      static final long REMOTE_FETCH_MAX_WAIT_MS = 6000L;
  
      private MockTime time;
      private ReplicaManager mockReplicaManager;
      private BrokerTopicStats brokerTopicStats;
      private SharePartitionManager sharePartitionManager;
  
      private static final List<TopicIdPartition> EMPTY_PART_LIST = List.of();
      private static final List<ShareFetchResponseData.AcquiredRecords> EMPTY_ACQUIRED_RECORDS = List.of();
  
      @BeforeEach
      public void setUp() {
          time = new MockTime();
          kafka.utils.TestUtils.clearYammerMetrics();
          brokerTopicStats = new BrokerTopicStats();
          mockReplicaManager = mock(ReplicaManager.class);
          Partition partition = mockPartition();
          when(mockReplicaManager.getPartitionOrException((TopicPartition) any())).thenReturn(partition);
      }
  
      @AfterEach
      public void tearDown() throws Exception {
          if (sharePartitionManager != null) {
              sharePartitionManager.close();
          }
          brokerTopicStats.close();
          assertNoReaperThreadsPendingClose();
      }
  
      @Test
      public void testNewContextReturnsFinalContextWithoutRequestData() {
          ShareSessionCache cache = new ShareSessionCache(10);
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withCache(cache)
              .build();
  
          Uuid tpId0 = Uuid.randomUuid();
          TopicIdPartition tp0 = new TopicIdPartition(tpId0, new TopicPartition("foo", 0));
          TopicIdPartition tp1 = new TopicIdPartition(tpId0, new TopicPartition("foo", 1));
  
          String groupId = "grp";
          Uuid memberId = Uuid.randomUuid();
  
          // Create a new share session with an initial share fetch request
          List<TopicIdPartition> reqData1 = List.of(tp0, tp1);
  
          ShareRequestMetadata reqMetadata1 = new ShareRequestMetadata(memberId, ShareRequestMetadata.INITIAL_EPOCH);
          ShareFetchContext context1 = sharePartitionManager.newContext(groupId, reqData1, EMPTY_PART_LIST, reqMetadata1, false, CONNECTION_ID);
          assertInstanceOf(ShareSessionContext.class, context1);
          assertFalse(((ShareSessionContext) context1).isSubsequent());
  
          ShareRequestMetadata reqMetadata2 = new ShareRequestMetadata(memberId, ShareRequestMetadata.FINAL_EPOCH);
          ShareFetchContext context2 = sharePartitionManager.newContext(groupId, EMPTY_PART_LIST, EMPTY_PART_LIST, reqMetadata2, true, CONNECTION_ID);
          assertEquals(FinalContext.class, context2.getClass());
      }
  
      @Test
      public void testNewContextReturnsFinalContextWithRequestData() {
          ShareSessionCache cache = new ShareSessionCache(10);
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withCache(cache)
              .build();
  
          Uuid tpId0 = Uuid.randomUuid();
          TopicIdPartition tp0 = new TopicIdPartition(tpId0, new TopicPartition("foo", 0));
          TopicIdPartition tp1 = new TopicIdPartition(tpId0, new TopicPartition("foo", 1));
  
          String groupId = "grp";
          Uuid memberId = Uuid.randomUuid();
  
          // Create a new share session with an initial share fetch request
          List<TopicIdPartition> reqData1 = List.of(tp0, tp1);
  
          ShareRequestMetadata reqMetadata1 = new ShareRequestMetadata(memberId, ShareRequestMetadata.INITIAL_EPOCH);
          ShareFetchContext context1 = sharePartitionManager.newContext(groupId, reqData1, EMPTY_PART_LIST, reqMetadata1, false, CONNECTION_ID);
          assertInstanceOf(ShareSessionContext.class, context1);
          assertFalse(((ShareSessionContext) context1).isSubsequent());
  
          ShareRequestMetadata reqMetadata2 = new ShareRequestMetadata(memberId, ShareRequestMetadata.FINAL_EPOCH);
  
          // Sending a Request with FINAL_EPOCH. This should return a FinalContext.
          List<TopicIdPartition> reqData2 = List.of(tp0, tp1);
          ShareFetchContext context2 = sharePartitionManager.newContext(groupId, reqData2, EMPTY_PART_LIST, reqMetadata2, true, CONNECTION_ID);
          assertEquals(FinalContext.class, context2.getClass());
      }
  
      @Test
      public void testNewContextReturnsFinalContextWhenTopicPartitionsArePresentInRequestData() {
          ShareSessionCache cache = new ShareSessionCache(10);
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withCache(cache)
              .build();
  
          Uuid tpId0 = Uuid.randomUuid();
          Uuid tpId1 = Uuid.randomUuid();
          TopicIdPartition tp0 = new TopicIdPartition(tpId0, new TopicPartition("foo", 0));
          TopicIdPartition tp1 = new TopicIdPartition(tpId0, new TopicPartition("foo", 1));
  
          String groupId = "grp";
          Uuid memberId = Uuid.randomUuid();
  
          // Create a new share session with an initial share fetch request
          List<TopicIdPartition> reqData1 = List.of(tp0, tp1);
  
          ShareRequestMetadata reqMetadata1 = new ShareRequestMetadata(memberId, ShareRequestMetadata.INITIAL_EPOCH);
          ShareFetchContext context1 = sharePartitionManager.newContext(groupId, reqData1, EMPTY_PART_LIST, reqMetadata1, false, CONNECTION_ID);
          assertInstanceOf(ShareSessionContext.class, context1);
          assertFalse(((ShareSessionContext) context1).isSubsequent());
  
          ShareRequestMetadata reqMetadata2 = new ShareRequestMetadata(memberId, ShareRequestMetadata.FINAL_EPOCH);
  
          // shareFetch is not empty, and it contains tpId1, which should return FinalContext instance since it is FINAL_EPOCH
          List<TopicIdPartition> reqData2 = List.of(new TopicIdPartition(tpId1, new TopicPartition("foo", 0)));
          assertInstanceOf(FinalContext.class,
              sharePartitionManager.newContext(groupId, reqData2, EMPTY_PART_LIST, reqMetadata2, true, CONNECTION_ID));
      }
  
      @Test
      public void testNewContextThrowsErrorWhenShareSessionNotFoundOnFinalEpoch() {
          ShareSessionCache cache = new ShareSessionCache(10);
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withCache(cache)
              .build();
          assertThrows(ShareSessionNotFoundException.class, () -> sharePartitionManager.newContext("grp", EMPTY_PART_LIST, EMPTY_PART_LIST,
              new ShareRequestMetadata(Uuid.randomUuid(), ShareRequestMetadata.FINAL_EPOCH), false, CONNECTION_ID));
      }
  
      @Test
      public void testNewContextThrowsErrorWhenAcknowledgeDataPresentOnInitialEpoch() {
          ShareSessionCache cache = new ShareSessionCache(10);
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withCache(cache)
              .build();
          Uuid tpId0 = Uuid.randomUuid();
          TopicIdPartition tp0 = new TopicIdPartition(tpId0, new TopicPartition("foo", 0));
          TopicIdPartition tp1 = new TopicIdPartition(tpId0, new TopicPartition("foo", 1));
  
          assertThrows(InvalidRequestException.class, () -> sharePartitionManager.newContext("grp", List.of(tp0, tp1), EMPTY_PART_LIST,
              new ShareRequestMetadata(Uuid.randomUuid(), ShareRequestMetadata.INITIAL_EPOCH), true, CONNECTION_ID));
      }
  
      @Test
      public void testNewContextThrowsErrorWhenShareSessionCacheIsFullOnInitialEpoch() {
          // Define a cache with max size 1
          ShareSessionCache cache = new ShareSessionCache(1);
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withCache(cache)
              .build();
  
          Uuid tpId0 = Uuid.randomUuid();
          TopicIdPartition tp0 = new TopicIdPartition(tpId0, new TopicPartition("foo", 0));
          TopicIdPartition tp1 = new TopicIdPartition(tpId0, new TopicPartition("foo", 1));
  
          String groupId = "grp";
          Uuid memberId1 = Uuid.randomUuid();
          Uuid memberId2 = Uuid.randomUuid();
  
          // Create a new share session with an initial share fetch request
          List<TopicIdPartition> reqData = List.of(tp0, tp1);
  
          ShareRequestMetadata reqMetadata1 = new ShareRequestMetadata(memberId1, ShareRequestMetadata.INITIAL_EPOCH);
          ShareFetchContext context1 = sharePartitionManager.newContext(groupId, reqData, EMPTY_PART_LIST, reqMetadata1, false, CONNECTION_ID);
          assertInstanceOf(ShareSessionContext.class, context1);
          assertFalse(((ShareSessionContext) context1).isSubsequent());
  
          // Trying to create a new share session, but since cache is already full, it should throw an exception
          ShareRequestMetadata reqMetadata2 = new ShareRequestMetadata(memberId2, ShareRequestMetadata.INITIAL_EPOCH);
          assertThrows(ShareSessionLimitReachedException.class, () -> sharePartitionManager.newContext("grp", reqData, EMPTY_PART_LIST,
              reqMetadata2, false, "id-2"));
      }
  
      @Test
      public void testNewContextExistingSessionNewRequestWithInitialEpoch() {
          ShareSessionCache cache = new ShareSessionCache(10);
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withCache(cache)
              .build();
  
          Uuid tpId0 = Uuid.randomUuid();
          TopicIdPartition tp0 = new TopicIdPartition(tpId0, new TopicPartition("foo", 0));
          TopicIdPartition tp1 = new TopicIdPartition(tpId0, new TopicPartition("foo", 1));
  
          String groupId = "grp";
          Uuid memberId = Uuid.randomUuid();
          List<TopicIdPartition> reqData = List.of(tp0, tp1);
  
          ShareRequestMetadata reqMetadata = new ShareRequestMetadata(memberId, ShareRequestMetadata.INITIAL_EPOCH);
  
          // Create a new share session with an initial share fetch request
          ShareFetchContext context1 = sharePartitionManager.newContext(groupId, reqData, EMPTY_PART_LIST, reqMetadata, false, CONNECTION_ID);
          assertInstanceOf(ShareSessionContext.class, context1);
          assertFalse(((ShareSessionContext) context1).isSubsequent());
          assertEquals(1, cache.size());
  
          // Sending another request with INITIAL_EPOCH and same share session key. This should return a new ShareSessionContext
          // and delete the older one.
          ShareFetchContext context2 = sharePartitionManager.newContext(groupId, reqData, EMPTY_PART_LIST, reqMetadata, false, CONNECTION_ID);
          assertInstanceOf(ShareSessionContext.class, context2);
          assertFalse(((ShareSessionContext) context1).isSubsequent());
          assertEquals(1, cache.size());
      }
  
      @Test
      public void testNewContext() {
          ShareSessionCache cache = new ShareSessionCache(10);
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withCache(cache)
              .build();
  
          Map<Uuid, String> topicNames = new HashMap<>();
          Uuid tpId0 = Uuid.randomUuid();
          Uuid tpId1 = Uuid.randomUuid();
          topicNames.put(tpId0, "foo");
          topicNames.put(tpId1, "bar");
          TopicIdPartition tp0 = new TopicIdPartition(tpId0, new TopicPartition("foo", 0));
          TopicIdPartition tp1 = new TopicIdPartition(tpId0, new TopicPartition("foo", 1));
          TopicIdPartition tp2 = new TopicIdPartition(tpId1, new TopicPartition("bar", 0));
          TopicIdPartition tp3 = new TopicIdPartition(tpId1, new TopicPartition("bar", 1));
  
          String groupId = "grp";
  
          // Create a new share session with an initial share fetch request
          List<TopicIdPartition> reqData2 = List.of(tp0, tp1);
  
          ShareRequestMetadata reqMetadata2 = new ShareRequestMetadata(Uuid.randomUuid(), ShareRequestMetadata.INITIAL_EPOCH);
          ShareFetchContext context2 = sharePartitionManager.newContext(groupId, reqData2, EMPTY_PART_LIST, reqMetadata2, false, CONNECTION_ID);
          assertInstanceOf(ShareSessionContext.class, context2);
          assertFalse(((ShareSessionContext) context2).isSubsequent());
  
          ((ShareSessionContext) context2).shareFetchData().forEach(topicIdPartition -> assertTrue(reqData2.contains(topicIdPartition)));
  
          LinkedHashMap<TopicIdPartition, ShareFetchResponseData.PartitionData> respData2 = new LinkedHashMap<>();
          respData2.put(tp0, new ShareFetchResponseData.PartitionData().setPartitionIndex(0));
          respData2.put(tp1, new ShareFetchResponseData.PartitionData().setPartitionIndex(1));
  
          ShareFetchResponse resp2 = context2.updateAndGenerateResponseData(groupId, reqMetadata2.memberId(), respData2);
          assertEquals(Errors.NONE, resp2.error());
          assertEquals(respData2, resp2.responseData(topicNames));
  
          ShareSessionKey shareSessionKey2 = new ShareSessionKey(groupId,
              reqMetadata2.memberId());
  
          // Test trying to create a new session with an invalid epoch
          assertThrows(InvalidShareSessionEpochException.class, () -> sharePartitionManager.newContext(groupId, reqData2, EMPTY_PART_LIST,
              new ShareRequestMetadata(shareSessionKey2.memberId(), 5), true, "id-2"));
  
          // Test trying to create a new session with a non-existent session key
          Uuid memberId4 = Uuid.randomUuid();
          assertThrows(ShareSessionNotFoundException.class, () -> sharePartitionManager.newContext(groupId, reqData2, EMPTY_PART_LIST,
              new ShareRequestMetadata(memberId4, 1), true, "id-3"));
  
          // Continue the first share session we created.
          ShareFetchContext context5 = sharePartitionManager.newContext(groupId, EMPTY_PART_LIST, EMPTY_PART_LIST,
              new ShareRequestMetadata(shareSessionKey2.memberId(), 1), true, CONNECTION_ID);
          assertInstanceOf(ShareSessionContext.class, context5);
          assertTrue(((ShareSessionContext) context5).isSubsequent());
  
          ShareSessionContext shareSessionContext5 = (ShareSessionContext) context5;
          synchronized (shareSessionContext5.session()) {
              shareSessionContext5.session().partitionMap().forEach(cachedSharePartition -> {
                  TopicIdPartition topicIdPartition = new TopicIdPartition(cachedSharePartition.topicId(), new
                      TopicPartition(cachedSharePartition.topic(), cachedSharePartition.partition()));
                  assertTrue(reqData2.contains(topicIdPartition));
              });
          }
          ShareFetchResponse resp5 = context5.updateAndGenerateResponseData(groupId, reqMetadata2.memberId(), respData2);
          assertEquals(Errors.NONE, resp5.error());
          assertEquals(0, resp5.responseData(topicNames).size());
  
          // Test setting an invalid share session epoch.
          assertThrows(InvalidShareSessionEpochException.class, () -> sharePartitionManager.newContext(groupId, reqData2, EMPTY_PART_LIST,
              new ShareRequestMetadata(shareSessionKey2.memberId(), 5), true, CONNECTION_ID));
  
          // Test generating a throttled response for a subsequent share session
          ShareFetchContext context7 = sharePartitionManager.newContext(groupId, EMPTY_PART_LIST, EMPTY_PART_LIST,
              new ShareRequestMetadata(shareSessionKey2.memberId(), 2), true, CONNECTION_ID);
          ShareFetchResponse resp7 = context7.throttleResponse(100);
          assertEquals(Errors.NONE, resp7.error());
          assertEquals(100, resp7.throttleTimeMs());
  
          // Get the final share session.
          ShareFetchContext context8 = sharePartitionManager.newContext(groupId, EMPTY_PART_LIST, EMPTY_PART_LIST,
              new ShareRequestMetadata(reqMetadata2.memberId(), ShareRequestMetadata.FINAL_EPOCH), true, CONNECTION_ID);
          assertEquals(FinalContext.class, context8.getClass());
          assertEquals(1, cache.size());
  
          LinkedHashMap<TopicIdPartition, ShareFetchResponseData.PartitionData> respData8 = new LinkedHashMap<>();
          respData8.put(tp2, new ShareFetchResponseData.PartitionData().setPartitionIndex(0));
          respData8.put(tp3, new ShareFetchResponseData.PartitionData().setPartitionIndex(1));
  
          ShareFetchResponse resp8 = context8.updateAndGenerateResponseData(groupId, reqMetadata2.memberId(), respData8);
          assertEquals(Errors.NONE, resp8.error());
  
          // Close the session.
          CompletableFuture<Map<TopicIdPartition, ShareAcknowledgeResponseData.PartitionData>> releaseResponse =
              sharePartitionManager.releaseSession(groupId, reqMetadata2.memberId().toString());
          assertTrue(releaseResponse.isDone());
          assertFalse(releaseResponse.isCompletedExceptionally());
          assertEquals(0, cache.size());
      }
  
      @Test
      public void testAcknowledgeSessionUpdateThrowsOnInitialEpoch() {
          ShareSessionCache cache = new ShareSessionCache(10);
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withCache(cache)
              .build();
  
          assertThrows(InvalidShareSessionEpochException.class,
              () -> sharePartitionManager.acknowledgeSessionUpdate("grp",
                  new ShareRequestMetadata(Uuid.randomUuid(), ShareRequestMetadata.INITIAL_EPOCH)));
      }
  
      @Test
      public void testAcknowledgeSessionUpdateThrowsWhenShareSessionNotFound() {
          ShareSessionCache cache = new ShareSessionCache(10);
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withCache(cache)
              .build();
  
          // The share session corresponding to this memberId has not been created yet. This should throw an exception.
          assertThrows(ShareSessionNotFoundException.class,
              () -> sharePartitionManager.acknowledgeSessionUpdate("grp",
                  new ShareRequestMetadata(Uuid.randomUuid(), ShareRequestMetadata.nextEpoch(ShareRequestMetadata.INITIAL_EPOCH))));
      }
  
      @Test
      public void testAcknowledgeSessionUpdateThrowsInvalidShareSessionEpochException() {
          ShareSessionCache cache = new ShareSessionCache(10);
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withCache(cache)
              .build();
  
          Uuid tpId0 = Uuid.randomUuid();
          TopicIdPartition tp0 = new TopicIdPartition(tpId0, new TopicPartition("foo", 0));
          TopicIdPartition tp1 = new TopicIdPartition(tpId0, new TopicPartition("foo", 1));
  
          String groupId = "grp";
          Uuid memberId = Uuid.randomUuid();
  
          // Create a new share session with an initial share fetch request
          ShareFetchContext context1 = sharePartitionManager.newContext(groupId, List.of(tp0, tp1), EMPTY_PART_LIST,
              new ShareRequestMetadata(memberId, ShareRequestMetadata.INITIAL_EPOCH), false, CONNECTION_ID);
          assertInstanceOf(ShareSessionContext.class, context1);
          assertFalse(((ShareSessionContext) context1).isSubsequent());
  
          // The expected epoch from the share session should be 1, but we are passing 2. This should throw an exception.
          assertThrows(InvalidShareSessionEpochException.class,
              () -> sharePartitionManager.acknowledgeSessionUpdate("grp",
                  new ShareRequestMetadata(memberId,
                      ShareRequestMetadata.nextEpoch(ShareRequestMetadata.nextEpoch(ShareRequestMetadata.INITIAL_EPOCH)))));
      }
  
      @Test
      public void testAcknowledgeSessionUpdateSuccessOnSubsequentEpoch() {
          ShareSessionCache cache = new ShareSessionCache(10);
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withCache(cache)
              .build();
  
          Uuid tpId0 = Uuid.randomUuid();
          TopicIdPartition tp0 = new TopicIdPartition(tpId0, new TopicPartition("foo", 0));
          TopicIdPartition tp1 = new TopicIdPartition(tpId0, new TopicPartition("foo", 1));
  
          String groupId = "grp";
          Uuid memberId = Uuid.randomUuid();
  
          // Create a new share session with an initial share fetch request
          ShareFetchContext context1 = sharePartitionManager.newContext(groupId, List.of(tp0, tp1), EMPTY_PART_LIST,
              new ShareRequestMetadata(memberId, ShareRequestMetadata.INITIAL_EPOCH), false, CONNECTION_ID);
          assertInstanceOf(ShareSessionContext.class, context1);
          assertFalse(((ShareSessionContext) context1).isSubsequent());
  
          // The expected epoch from the share session should be 1, and we are passing the same. So, execution should be successful.
          assertDoesNotThrow(
              () -> sharePartitionManager.acknowledgeSessionUpdate("grp",
                  new ShareRequestMetadata(memberId, ShareRequestMetadata.nextEpoch(ShareRequestMetadata.INITIAL_EPOCH))));
      }
  
      @Test
      public void testAcknowledgeSessionUpdateSuccessOnFinalEpoch() {
          ShareSessionCache cache = new ShareSessionCache(10);
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withCache(cache)
              .build();
  
          Uuid tpId0 = Uuid.randomUuid();
          TopicIdPartition tp0 = new TopicIdPartition(tpId0, new TopicPartition("foo", 0));
          TopicIdPartition tp1 = new TopicIdPartition(tpId0, new TopicPartition("foo", 1));
  
          String groupId = "grp";
          Uuid memberId = Uuid.randomUuid();
  
          // Create a new share session with an initial share fetch request
          ShareFetchContext context1 = sharePartitionManager.newContext(groupId, List.of(tp0, tp1), EMPTY_PART_LIST,
              new ShareRequestMetadata(memberId, ShareRequestMetadata.INITIAL_EPOCH), false, CONNECTION_ID);
          assertInstanceOf(ShareSessionContext.class, context1);
          assertFalse(((ShareSessionContext) context1).isSubsequent());
  
          // The expected epoch from the share session should be 1, but we are passing the Final Epoch (-1). This should throw an exception.
          assertDoesNotThrow(
              () -> sharePartitionManager.acknowledgeSessionUpdate("grp",
                  new ShareRequestMetadata(memberId, ShareRequestMetadata.FINAL_EPOCH)));
      }
  
      @Test
      public void testSubsequentShareSession() {
          sharePartitionManager = SharePartitionManagerBuilder.builder().build();
          Map<Uuid, String> topicNames = new HashMap<>();
          Uuid fooId = Uuid.randomUuid();
          Uuid barId = Uuid.randomUuid();
          topicNames.put(fooId, "foo");
          topicNames.put(barId, "bar");
          TopicIdPartition tp0 = new TopicIdPartition(fooId, new TopicPartition("foo", 0));
          TopicIdPartition tp1 = new TopicIdPartition(fooId, new TopicPartition("foo", 1));
          TopicIdPartition tp2 = new TopicIdPartition(barId, new TopicPartition("bar", 0));
  
          // Create a new share session with foo-0 and foo-1
          List<TopicIdPartition> reqData1 = List.of(tp0, tp1);
  
          String groupId = "grp";
          ShareRequestMetadata reqMetadata1 = new ShareRequestMetadata(Uuid.randomUuid(), ShareRequestMetadata.INITIAL_EPOCH);
  
          ShareFetchContext context1 = sharePartitionManager.newContext(groupId, reqData1, EMPTY_PART_LIST, reqMetadata1, false, CONNECTION_ID);
          assertInstanceOf(ShareSessionContext.class, context1);
  
          LinkedHashMap<TopicIdPartition, ShareFetchResponseData.PartitionData> respData1 = new LinkedHashMap<>();
          respData1.put(tp0, new ShareFetchResponseData.PartitionData().setPartitionIndex(tp0.partition()));
          respData1.put(tp1, new ShareFetchResponseData.PartitionData().setPartitionIndex(tp1.partition()));
  
          ShareFetchResponse resp1 = context1.updateAndGenerateResponseData(groupId, reqMetadata1.memberId(), respData1);
          assertEquals(Errors.NONE, resp1.error());
          assertEquals(2, resp1.responseData(topicNames).size());
  
          // Create a subsequent fetch request that removes foo-0 and adds bar-0
          List<TopicIdPartition> reqData2 = List.of(tp2);
          List<TopicIdPartition> removed2 = new ArrayList<>();
          removed2.add(tp0);
          ShareFetchContext context2 = sharePartitionManager.newContext(groupId, reqData2, removed2,
                  new ShareRequestMetadata(reqMetadata1.memberId(), 1), true, CONNECTION_ID);
          assertInstanceOf(ShareSessionContext.class, context2);
  
          Set<TopicIdPartition> expectedTopicIdPartitions2 = new HashSet<>();
          expectedTopicIdPartitions2.add(tp1);
          expectedTopicIdPartitions2.add(tp2);
          Set<TopicIdPartition> actualTopicIdPartitions2 = new HashSet<>();
          ShareSessionContext shareSessionContext = (ShareSessionContext) context2;
          shareSessionContext.session().partitionMap().forEach(cachedSharePartition -> {
              TopicIdPartition topicIdPartition = new TopicIdPartition(cachedSharePartition.topicId(), new
                      TopicPartition(cachedSharePartition.topic(), cachedSharePartition.partition()));
              actualTopicIdPartitions2.add(topicIdPartition);
          });
  
          assertEquals(expectedTopicIdPartitions2, actualTopicIdPartitions2);
          LinkedHashMap<TopicIdPartition, ShareFetchResponseData.PartitionData> respData2 = new LinkedHashMap<>();
          respData2.put(tp1, new ShareFetchResponseData.PartitionData().setPartitionIndex(tp1.partition()));
          respData2.put(tp2, new ShareFetchResponseData.PartitionData().setPartitionIndex(tp2.partition()));
  
          ShareFetchResponse resp2 = context2.updateAndGenerateResponseData(groupId, reqMetadata1.memberId(), respData2);
          assertEquals(Errors.NONE, resp2.error());
          assertEquals(1, resp2.data().responses().size());
          assertEquals(barId, resp2.data().responses().stream().findFirst().get().topicId());
          assertEquals(1, resp2.data().responses().stream().findFirst().get().partitions().size());
          assertEquals(0, resp2.data().responses().stream().findFirst().get().partitions().get(0).partitionIndex());
          assertEquals(1, resp2.responseData(topicNames).size());
      }
  
      @Test
      public void testZeroSizeShareSession() {
          ShareSessionCache cache = new ShareSessionCache(10);
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withCache(cache)
              .build();
  
          Map<Uuid, String> topicNames = new HashMap<>();
          Uuid fooId = Uuid.randomUuid();
          topicNames.put(fooId, "foo");
          TopicIdPartition foo0 = new TopicIdPartition(fooId, new TopicPartition("foo", 0));
          TopicIdPartition foo1 = new TopicIdPartition(fooId, new TopicPartition("foo", 1));
  
          // Create a new share session with foo-0 and foo-1
          List<TopicIdPartition> reqData1 = List.of(foo0, foo1);
  
          String groupId = "grp";
          ShareRequestMetadata reqMetadata1 = new ShareRequestMetadata(Uuid.randomUuid(), ShareRequestMetadata.INITIAL_EPOCH);
  
          ShareFetchContext context1 = sharePartitionManager.newContext(groupId, reqData1, EMPTY_PART_LIST, reqMetadata1, false, CONNECTION_ID);
          assertInstanceOf(ShareSessionContext.class, context1);
  
          LinkedHashMap<TopicIdPartition, ShareFetchResponseData.PartitionData> respData1 = new LinkedHashMap<>();
          respData1.put(foo0, new ShareFetchResponseData.PartitionData().setPartitionIndex(foo0.partition()));
          respData1.put(foo1, new ShareFetchResponseData.PartitionData().setPartitionIndex(foo1.partition()));
  
          ShareFetchResponse resp1 = context1.updateAndGenerateResponseData(groupId, reqMetadata1.memberId(), respData1);
          assertEquals(Errors.NONE, resp1.error());
          assertEquals(2, resp1.responseData(topicNames).size());
  
          // Create a subsequent share request that removes foo-0 and foo-1
          // Verify that the previous share session was closed.
          List<TopicIdPartition> removed2 = new ArrayList<>();
          removed2.add(foo0);
          removed2.add(foo1);
          ShareFetchContext context2 = sharePartitionManager.newContext(groupId, EMPTY_PART_LIST, removed2,
                  new ShareRequestMetadata(reqMetadata1.memberId(), 1), true, CONNECTION_ID);
          assertInstanceOf(ShareSessionContext.class, context2);
  
          LinkedHashMap<TopicIdPartition, ShareFetchResponseData.PartitionData> respData2 = new LinkedHashMap<>();
          ShareFetchResponse resp2 = context2.updateAndGenerateResponseData(groupId, reqMetadata1.memberId(), respData2);
          assertTrue(resp2.responseData(topicNames).isEmpty());
          assertEquals(1, cache.size());
      }
  
      @Test
      public void testToForgetPartitions() {
          String groupId = "grp";
          ShareSessionCache cache = new ShareSessionCache(10);
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withCache(cache)
              .build();
  
          Uuid fooId = Uuid.randomUuid();
          Uuid barId = Uuid.randomUuid();
          TopicIdPartition foo = new TopicIdPartition(fooId, new TopicPartition("foo", 0));
          TopicIdPartition bar = new TopicIdPartition(barId, new TopicPartition("bar", 0));
  
          ShareRequestMetadata reqMetadata1 = new ShareRequestMetadata(Uuid.randomUuid(), ShareRequestMetadata.INITIAL_EPOCH);
  
          List<TopicIdPartition> reqData1 = List.of(foo, bar);
  
          ShareFetchContext context1 = sharePartitionManager.newContext(groupId, reqData1, EMPTY_PART_LIST, reqMetadata1, false, CONNECTION_ID);
          assertInstanceOf(ShareSessionContext.class, context1);
          assertPartitionsPresent((ShareSessionContext) context1, List.of(foo, bar));
  
          mockUpdateAndGenerateResponseData(context1, groupId, reqMetadata1.memberId());
  
          ShareFetchContext context2 = sharePartitionManager.newContext(groupId, EMPTY_PART_LIST, List.of(foo),
                  new ShareRequestMetadata(reqMetadata1.memberId(), 1), true, CONNECTION_ID);
  
          // So foo is removed but not the others.
          assertPartitionsPresent((ShareSessionContext) context2, List.of(bar));
  
          mockUpdateAndGenerateResponseData(context2, groupId, reqMetadata1.memberId());
  
          ShareFetchContext context3 = sharePartitionManager.newContext(groupId, EMPTY_PART_LIST, List.of(bar),
                  new ShareRequestMetadata(reqMetadata1.memberId(), 2), true, CONNECTION_ID);
          assertPartitionsPresent((ShareSessionContext) context3, EMPTY_PART_LIST);
      }
  
      // This test simulates a share session where the topic ID changes broker side (the one handling the request) in both the metadata cache and the log
      // -- as though the topic is deleted and recreated.
      @Test
      public void testShareSessionUpdateTopicIdsBrokerSide() {
          String groupId = "grp";
          ShareSessionCache cache = new ShareSessionCache(10);
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withCache(cache)
              .build();
  
          Uuid fooId = Uuid.randomUuid();
          Uuid barId = Uuid.randomUuid();
          TopicIdPartition foo = new TopicIdPartition(fooId, new TopicPartition("foo", 0));
          TopicIdPartition bar = new TopicIdPartition(barId, new TopicPartition("bar", 1));
  
          Map<Uuid, String> topicNames = new HashMap<>();
          topicNames.put(fooId, "foo");
          topicNames.put(barId, "bar");
  
          // Create a new share session with foo-0 and bar-1
          List<TopicIdPartition> reqData1 = List.of(foo, bar);
  
          ShareRequestMetadata reqMetadata1 = new ShareRequestMetadata(Uuid.randomUuid(), ShareRequestMetadata.INITIAL_EPOCH);
          ShareFetchContext context1 = sharePartitionManager.newContext(groupId, reqData1, EMPTY_PART_LIST, reqMetadata1, false, CONNECTION_ID);
  
          assertInstanceOf(ShareSessionContext.class, context1);
          assertFalse(((ShareSessionContext) context1).isSubsequent());
  
          LinkedHashMap<TopicIdPartition, ShareFetchResponseData.PartitionData> respData1 = new LinkedHashMap<>();
          respData1.put(bar, new ShareFetchResponseData.PartitionData().setPartitionIndex(bar.partition()));
          respData1.put(foo, new ShareFetchResponseData.PartitionData().setPartitionIndex(foo.partition()).setErrorCode(
                  Errors.UNKNOWN_TOPIC_OR_PARTITION.code()));
  
          ShareFetchResponse resp1 = context1.updateAndGenerateResponseData(groupId, reqMetadata1.memberId(), respData1);
          assertEquals(Errors.NONE, resp1.error());
          assertEquals(2, resp1.responseData(topicNames).size());
  
          // Create a subsequent share fetch request as though no topics changed.
          ShareFetchContext context2 = sharePartitionManager.newContext(groupId, EMPTY_PART_LIST, EMPTY_PART_LIST,
                  new ShareRequestMetadata(reqMetadata1.memberId(), 1), true, CONNECTION_ID);
  
          assertInstanceOf(ShareSessionContext.class, context2);
          assertTrue(((ShareSessionContext) context2).isSubsequent());
  
          LinkedHashMap<TopicIdPartition, ShareFetchResponseData.PartitionData> respData2 = new LinkedHashMap<>();
          // Likely if the topic ID is different in the broker, it will be different in the log. Simulate the log check finding an inconsistent ID.
          respData2.put(foo, new ShareFetchResponseData.PartitionData().setPartitionIndex(foo.partition()).setErrorCode(
                  Errors.INCONSISTENT_TOPIC_ID.code()));
          ShareFetchResponse resp2 = context2.updateAndGenerateResponseData(groupId, reqMetadata1.memberId(), respData2);
          assertEquals(Errors.NONE, resp2.error());
          // We should have the inconsistent topic ID error on the partition
          assertEquals(Errors.INCONSISTENT_TOPIC_ID.code(), resp2.responseData(topicNames).get(foo).errorCode());
      }
  
      @Test
      public void testGetErroneousAndValidTopicIdPartitions() {
          ShareSessionCache cache = new ShareSessionCache(10);
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withCache(cache)
              .build();
  
          Uuid tpId0 = Uuid.randomUuid();
          TopicIdPartition tp0 = new TopicIdPartition(tpId0, new TopicPartition("foo", 0));
          TopicIdPartition tp1 = new TopicIdPartition(tpId0, new TopicPartition("foo", 1));
          TopicIdPartition tpNull1 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition(null, 0));
          TopicIdPartition tpNull2 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition(null, 1));
  
          String groupId = "grp";
  
          // Create a new share session with an initial share fetch request
          List<TopicIdPartition> reqData2 = List.of(tp0, tp1, tpNull1);
  
          ShareRequestMetadata reqMetadata2 = new ShareRequestMetadata(Uuid.randomUuid(), ShareRequestMetadata.INITIAL_EPOCH);
          ShareFetchContext context2 = sharePartitionManager.newContext(groupId, reqData2, EMPTY_PART_LIST, reqMetadata2, false, CONNECTION_ID);
          assertInstanceOf(ShareSessionContext.class, context2);
          assertFalse(((ShareSessionContext) context2).isSubsequent());
          assertErroneousAndValidTopicIdPartitions(context2.getErroneousAndValidTopicIdPartitions(), List.of(tpNull1), List.of(tp0, tp1));
  
          LinkedHashMap<TopicIdPartition, ShareFetchResponseData.PartitionData> respData2 = new LinkedHashMap<>();
          respData2.put(tp0, new ShareFetchResponseData.PartitionData().setPartitionIndex(0));
          respData2.put(tp1, new ShareFetchResponseData.PartitionData().setPartitionIndex(1));
          respData2.put(tpNull1, new ShareFetchResponseData.PartitionData().setPartitionIndex(0));
  
          ShareFetchResponse resp2 = context2.updateAndGenerateResponseData(groupId, reqMetadata2.memberId(), respData2);
          assertEquals(Errors.NONE, resp2.error());
  
          ShareSessionKey shareSessionKey2 = new ShareSessionKey(groupId, reqMetadata2.memberId());
  
          // Check for throttled response
          ShareFetchResponse resp2Throttle = context2.throttleResponse(100);
          assertEquals(Errors.NONE, resp2Throttle.error());
          assertEquals(100, resp2Throttle.throttleTimeMs());
  
          // Test trying to create a new session with an invalid epoch
          assertThrows(InvalidShareSessionEpochException.class, () -> sharePartitionManager.newContext(groupId, reqData2, EMPTY_PART_LIST,
                  new ShareRequestMetadata(shareSessionKey2.memberId(), 5), true, CONNECTION_ID));
  
          // Test trying to create a new session with a non-existent session key
          assertThrows(ShareSessionNotFoundException.class, () -> sharePartitionManager.newContext(groupId, reqData2, EMPTY_PART_LIST,
                  new ShareRequestMetadata(Uuid.randomUuid(), 1), true, CONNECTION_ID));
  
          // Continue the first share session we created.
          ShareFetchContext context5 = sharePartitionManager.newContext(groupId, EMPTY_PART_LIST, EMPTY_PART_LIST,
                  new ShareRequestMetadata(shareSessionKey2.memberId(), 1), true, CONNECTION_ID);
          assertInstanceOf(ShareSessionContext.class, context5);
          assertTrue(((ShareSessionContext) context5).isSubsequent());
  
          assertErroneousAndValidTopicIdPartitions(context5.getErroneousAndValidTopicIdPartitions(), List.of(tpNull1), List.of(tp0, tp1));
  
          ShareFetchResponse resp5 = context5.updateAndGenerateResponseData(groupId, reqMetadata2.memberId(), respData2);
          assertEquals(Errors.NONE, resp5.error());
  
          // Test setting an invalid share session epoch.
          assertThrows(InvalidShareSessionEpochException.class, () -> sharePartitionManager.newContext(groupId, reqData2, EMPTY_PART_LIST,
                  new ShareRequestMetadata(shareSessionKey2.memberId(), 5), true, CONNECTION_ID));
  
          // Test generating a throttled response for a subsequent share session
          List<TopicIdPartition> reqData7 = List.of(tpNull2);
          ShareFetchContext context7 = sharePartitionManager.newContext(groupId, reqData7, EMPTY_PART_LIST,
                  new ShareRequestMetadata(shareSessionKey2.memberId(), 2), true, CONNECTION_ID);
          // Check for throttled response
          ShareFetchResponse resp7 = context7.throttleResponse(100);
          assertEquals(Errors.NONE, resp7.error());
          assertEquals(100, resp7.throttleTimeMs());
  
          assertErroneousAndValidTopicIdPartitions(context7.getErroneousAndValidTopicIdPartitions(), List.of(tpNull1, tpNull2), List.of(tp0, tp1));
  
          // Get the final share session.
          ShareFetchContext context8 = sharePartitionManager.newContext(groupId, EMPTY_PART_LIST, EMPTY_PART_LIST,
                  new ShareRequestMetadata(reqMetadata2.memberId(), ShareRequestMetadata.FINAL_EPOCH), true, CONNECTION_ID);
          assertEquals(FinalContext.class, context8.getClass());
          assertEquals(1, cache.size());
  
          assertErroneousAndValidTopicIdPartitions(context8.getErroneousAndValidTopicIdPartitions(), EMPTY_PART_LIST, EMPTY_PART_LIST);
          // Check for throttled response
          ShareFetchResponse resp8 = context8.throttleResponse(100);
          assertEquals(Errors.NONE, resp8.error());
          assertEquals(100, resp8.throttleTimeMs());
  
          // Close the session.
          CompletableFuture<Map<TopicIdPartition, ShareAcknowledgeResponseData.PartitionData>> releaseResponse =
              sharePartitionManager.releaseSession(groupId, reqMetadata2.memberId().toString());
          assertTrue(releaseResponse.isDone());
          assertFalse(releaseResponse.isCompletedExceptionally());
          assertEquals(0, cache.size());
      }
  
      @Test
      public void testShareFetchContextResponseSize() {
          ShareSessionCache cache = new ShareSessionCache(10);
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withCache(cache)
              .build();
  
          Map<Uuid, String> topicNames = new HashMap<>();
          Uuid tpId0 = Uuid.randomUuid();
          Uuid tpId1 = Uuid.randomUuid();
          topicNames.put(tpId0, "foo");
          topicNames.put(tpId1, "bar");
          TopicIdPartition tp0 = new TopicIdPartition(tpId0, new TopicPartition("foo", 0));
          TopicIdPartition tp1 = new TopicIdPartition(tpId0, new TopicPartition("foo", 1));
          TopicIdPartition tp2 = new TopicIdPartition(tpId1, new TopicPartition("bar", 0));
          TopicIdPartition tp3 = new TopicIdPartition(tpId1, new TopicPartition("bar", 1));
  
          String groupId = "grp";
  
          // Create a new share session with an initial share fetch request
          List<TopicIdPartition> reqData2 = List.of(tp0, tp1);
  
          // For response size expected value calculation
          ObjectSerializationCache objectSerializationCache = new ObjectSerializationCache();
          short version = ApiKeys.SHARE_FETCH.latestVersion();
  
          ShareRequestMetadata reqMetadata2 = new ShareRequestMetadata(Uuid.randomUuid(), ShareRequestMetadata.INITIAL_EPOCH);
          ShareFetchContext context2 = sharePartitionManager.newContext(groupId, reqData2, EMPTY_PART_LIST, reqMetadata2, false, CONNECTION_ID);
          assertInstanceOf(ShareSessionContext.class, context2);
          assertFalse(((ShareSessionContext) context2).isSubsequent());
  
          LinkedHashMap<TopicIdPartition, ShareFetchResponseData.PartitionData> respData2 = new LinkedHashMap<>();
          respData2.put(tp0, new ShareFetchResponseData.PartitionData().setPartitionIndex(0));
          respData2.put(tp1, new ShareFetchResponseData.PartitionData().setPartitionIndex(1));
  
          int respSize2 = context2.responseSize(respData2, version);
          ShareFetchResponse resp2 = context2.updateAndGenerateResponseData(groupId, reqMetadata2.memberId(), respData2);
          assertEquals(Errors.NONE, resp2.error());
          assertEquals(respData2, resp2.responseData(topicNames));
          // We add 4 here in response to 4 being added in sizeOf() method in ShareFetchResponse class.
          assertEquals(4 + resp2.data().size(objectSerializationCache, version), respSize2);
  
          ShareSessionKey shareSessionKey2 = new ShareSessionKey(groupId,
                  reqMetadata2.memberId());
  
          // Test trying to create a new session with an invalid epoch
          assertThrows(InvalidShareSessionEpochException.class, () -> sharePartitionManager.newContext(groupId, reqData2, EMPTY_PART_LIST,
                  new ShareRequestMetadata(shareSessionKey2.memberId(), 5), true, CONNECTION_ID));
  
          // Test trying to create a new session with a non-existent session key
          Uuid memberId4 = Uuid.randomUuid();
          assertThrows(ShareSessionNotFoundException.class, () -> sharePartitionManager.newContext(groupId, reqData2, EMPTY_PART_LIST,
                  new ShareRequestMetadata(memberId4, 1), true, CONNECTION_ID));
  
          // Continue the first share session we created.
          List<TopicIdPartition> reqData5 = List.of(tp2);
          ShareFetchContext context5 = sharePartitionManager.newContext(groupId, reqData5, EMPTY_PART_LIST,
                  new ShareRequestMetadata(shareSessionKey2.memberId(), 1), true, CONNECTION_ID);
          assertInstanceOf(ShareSessionContext.class, context5);
          assertTrue(((ShareSessionContext) context5).isSubsequent());
  
          LinkedHashMap<TopicIdPartition, ShareFetchResponseData.PartitionData> respData5 = new LinkedHashMap<>();
          respData5.put(tp2, new ShareFetchResponseData.PartitionData().setPartitionIndex(0));
  
          int respSize5 = context5.responseSize(respData5, version);
          ShareFetchResponse resp5 = context5.updateAndGenerateResponseData(groupId, reqMetadata2.memberId(), respData5);
          assertEquals(Errors.NONE, resp5.error());
          // We add 4 here in response to 4 being added in sizeOf() method in ShareFetchResponse class.
          assertEquals(4 + resp5.data().size(objectSerializationCache, version), respSize5);
  
          // Test setting an invalid share session epoch.
          assertThrows(InvalidShareSessionEpochException.class, () -> sharePartitionManager.newContext(groupId, reqData2, EMPTY_PART_LIST,
                  new ShareRequestMetadata(shareSessionKey2.memberId(), 5), true, CONNECTION_ID));
  
          // Test generating a throttled response for a subsequent share session
          ShareFetchContext context7 = sharePartitionManager.newContext(groupId, EMPTY_PART_LIST, EMPTY_PART_LIST,
                  new ShareRequestMetadata(shareSessionKey2.memberId(), 2), true, CONNECTION_ID);
  
          int respSize7 = context7.responseSize(respData2, version);
          ShareFetchResponse resp7 = context7.throttleResponse(100);
          assertEquals(Errors.NONE, resp7.error());
          assertEquals(100, resp7.throttleTimeMs());
          // We add 4 here in response to 4 being added in sizeOf() method in ShareFetchResponse class.
          assertEquals(4 + new ShareFetchResponseData().size(objectSerializationCache, version), respSize7);
  
          // Get the final share session.
          ShareFetchContext context8 = sharePartitionManager.newContext(groupId, EMPTY_PART_LIST, EMPTY_PART_LIST,
                  new ShareRequestMetadata(reqMetadata2.memberId(), ShareRequestMetadata.FINAL_EPOCH), true, CONNECTION_ID);
          assertEquals(FinalContext.class, context8.getClass());
          assertEquals(1, cache.size());
  
          LinkedHashMap<TopicIdPartition, ShareFetchResponseData.PartitionData> respData8 = new LinkedHashMap<>();
          respData8.put(tp3, new ShareFetchResponseData.PartitionData().setPartitionIndex(1));
  
          int respSize8 = context8.responseSize(respData8, version);
          ShareFetchResponse resp8 = context8.updateAndGenerateResponseData(groupId, reqMetadata2.memberId(), respData8);
          assertEquals(Errors.NONE, resp8.error());
          // We add 4 here in response to 4 being added in sizeOf() method in ShareFetchResponse class.
          assertEquals(4 + resp8.data().size(objectSerializationCache, version), respSize8);
      }
  
      @Test
      public void testCachedTopicPartitionsWithNoTopicPartitions() {
          ShareSessionCache cache = new ShareSessionCache(10);
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withCache(cache)
              .build();
  
          List<TopicIdPartition> result = sharePartitionManager.cachedTopicIdPartitionsInShareSession("grp", Uuid.randomUuid());
          assertTrue(result.isEmpty());
      }
  
      @Test
      public void testCachedTopicPartitionsForValidShareSessions() {
          ShareSessionCache cache = new ShareSessionCache(10);
  
          Uuid tpId0 = Uuid.randomUuid();
          Uuid tpId1 = Uuid.randomUuid();
          TopicIdPartition tp0 = new TopicIdPartition(tpId0, new TopicPartition("foo", 0));
          TopicIdPartition tp1 = new TopicIdPartition(tpId0, new TopicPartition("foo", 1));
          TopicIdPartition tp2 = new TopicIdPartition(tpId1, new TopicPartition("bar", 0));
          TopicIdPartition tp3 = new TopicIdPartition(tpId1, new TopicPartition("bar", 1));
          String groupId = "grp";
          Uuid memberId1 = Uuid.randomUuid();
          Uuid memberId2 = Uuid.randomUuid();
          SharePartition sp0 = mock(SharePartition.class);
          SharePartition sp1 = mock(SharePartition.class);
          SharePartition sp2 = mock(SharePartition.class);
  
          when(sp0.releaseAcquiredRecords(ArgumentMatchers.eq(String.valueOf(memberId1)))).thenReturn(CompletableFuture.completedFuture(null));
          when(sp1.releaseAcquiredRecords(ArgumentMatchers.eq(String.valueOf(memberId1)))).thenReturn(CompletableFuture.completedFuture(null));
          when(sp2.releaseAcquiredRecords(ArgumentMatchers.eq(String.valueOf(memberId1)))).thenReturn(CompletableFuture.completedFuture(null));
  
          SharePartitionCache partitionCache = new SharePartitionCache();
          partitionCache.put(new SharePartitionKey(groupId, tp0), sp0);
          partitionCache.put(new SharePartitionKey(groupId, tp1), sp1);
          partitionCache.put(new SharePartitionKey(groupId, tp2), sp2);
  
          sharePartitionManager = SharePartitionManagerBuilder.builder()
                  .withCache(cache)
                  .withPartitionCache(partitionCache)
                  .build();
  
  
          // Create a new share session with an initial share fetch request.
          List<TopicIdPartition> reqData1 = List.of(tp0, tp1);
  
          ShareRequestMetadata reqMetadata1 = new ShareRequestMetadata(memberId1, ShareRequestMetadata.INITIAL_EPOCH);
          ShareFetchContext context1 = sharePartitionManager.newContext(groupId, reqData1, EMPTY_PART_LIST, reqMetadata1, false, CONNECTION_ID);
          assertInstanceOf(ShareSessionContext.class, context1);
          assertFalse(((ShareSessionContext) context1).isSubsequent());
  
          ShareSessionKey shareSessionKey1 = new ShareSessionKey(groupId,
                  reqMetadata1.memberId());
  
          LinkedHashMap<TopicIdPartition, ShareFetchResponseData.PartitionData> respData1 = new LinkedHashMap<>();
          respData1.put(tp0, new ShareFetchResponseData.PartitionData().setPartitionIndex(0));
          respData1.put(tp1, new ShareFetchResponseData.PartitionData().setPartitionIndex(1));
  
          ShareFetchResponse resp1 = context1.updateAndGenerateResponseData(groupId, reqMetadata1.memberId(), respData1);
          assertEquals(Errors.NONE, resp1.error());
  
          assertEquals(new HashSet<>(List.of(tp0, tp1)),
                  new HashSet<>(sharePartitionManager.cachedTopicIdPartitionsInShareSession(groupId, memberId1)));
  
          // Create a new share session with an initial share fetch request.
          List<TopicIdPartition> reqData2 = List.of(tp2);
  
          ShareRequestMetadata reqMetadata2 = new ShareRequestMetadata(memberId2, ShareRequestMetadata.INITIAL_EPOCH);
          ShareFetchContext context2 = sharePartitionManager.newContext(groupId, reqData2, EMPTY_PART_LIST, reqMetadata2, false, CONNECTION_ID);
          assertInstanceOf(ShareSessionContext.class, context2);
          assertFalse(((ShareSessionContext) context2).isSubsequent());
  
          ShareSessionKey shareSessionKey2 = new ShareSessionKey(groupId,
                  reqMetadata2.memberId());
  
          LinkedHashMap<TopicIdPartition, ShareFetchResponseData.PartitionData> respData2 = new LinkedHashMap<>();
          respData2.put(tp2, new ShareFetchResponseData.PartitionData().setPartitionIndex(0));
  
          ShareFetchResponse resp2 = context2.updateAndGenerateResponseData(groupId, reqMetadata2.memberId(), respData2);
          assertEquals(Errors.NONE, resp2.error());
  
          assertEquals(List.of(tp2), sharePartitionManager.cachedTopicIdPartitionsInShareSession(groupId, memberId2));
  
          // Continue the first share session we created.
          List<TopicIdPartition> reqData3 = List.of(tp2);
          ShareFetchContext context3 = sharePartitionManager.newContext(groupId, reqData3, EMPTY_PART_LIST,
                  new ShareRequestMetadata(shareSessionKey1.memberId(), 1), true, CONNECTION_ID);
          assertInstanceOf(ShareSessionContext.class, context3);
          assertTrue(((ShareSessionContext) context3).isSubsequent());
  
          LinkedHashMap<TopicIdPartition, ShareFetchResponseData.PartitionData> respData3 = new LinkedHashMap<>();
          respData3.put(tp2, new ShareFetchResponseData.PartitionData().setPartitionIndex(0));
          ShareFetchResponse resp3 = context3.updateAndGenerateResponseData(groupId, reqMetadata1.memberId(), respData3);
          assertEquals(Errors.NONE, resp3.error());
  
          assertEquals(new HashSet<>(List.of(tp0, tp1, tp2)),
                  new HashSet<>(sharePartitionManager.cachedTopicIdPartitionsInShareSession(groupId, memberId1)));
  
          // Continue the second session we created.
          List<TopicIdPartition> reqData4 = List.of(tp3);
          ShareFetchContext context4 = sharePartitionManager.newContext(groupId, reqData4, List.of(tp2),
                  new ShareRequestMetadata(shareSessionKey2.memberId(), 1), true, CONNECTION_ID);
          assertInstanceOf(ShareSessionContext.class, context4);
          assertTrue(((ShareSessionContext) context4).isSubsequent());
  
          LinkedHashMap<TopicIdPartition, ShareFetchResponseData.PartitionData> respData4 = new LinkedHashMap<>();
          respData4.put(tp3, new ShareFetchResponseData.PartitionData().setPartitionIndex(1));
          ShareFetchResponse resp4 = context4.updateAndGenerateResponseData(groupId, reqMetadata2.memberId(), respData4);
          assertEquals(Errors.NONE, resp4.error());
  
          assertEquals(List.of(tp3), sharePartitionManager.cachedTopicIdPartitionsInShareSession(groupId, memberId2));
  
          // Get the final share session.
          ShareFetchContext context5 = sharePartitionManager.newContext(groupId, EMPTY_PART_LIST, EMPTY_PART_LIST,
                  new ShareRequestMetadata(reqMetadata1.memberId(), ShareRequestMetadata.FINAL_EPOCH), true, CONNECTION_ID);
          assertEquals(FinalContext.class, context5.getClass());
  
          LinkedHashMap<TopicIdPartition, ShareFetchResponseData.PartitionData> respData5 = new LinkedHashMap<>();
          ShareFetchResponse resp5 = context5.updateAndGenerateResponseData(groupId, reqMetadata1.memberId(), respData5);
          assertEquals(Errors.NONE, resp5.error());
  
          assertFalse(sharePartitionManager.cachedTopicIdPartitionsInShareSession(groupId, memberId1).isEmpty());
  
          // Close the first session.
          sharePartitionManager.releaseSession(groupId, reqMetadata1.memberId().toString());
          assertTrue(sharePartitionManager.cachedTopicIdPartitionsInShareSession(groupId, memberId1).isEmpty());
  
          // Continue the second share session .
          ShareFetchContext context6 = sharePartitionManager.newContext(groupId, EMPTY_PART_LIST, List.of(tp3),
                  new ShareRequestMetadata(shareSessionKey2.memberId(), 2), true, CONNECTION_ID);
          assertInstanceOf(ShareSessionContext.class, context6);
          assertTrue(((ShareSessionContext) context6).isSubsequent());
  
          LinkedHashMap<TopicIdPartition, ShareFetchResponseData.PartitionData> respData6 = new LinkedHashMap<>();
          ShareFetchResponse resp6 = context6.updateAndGenerateResponseData(groupId, reqMetadata2.memberId(), respData6);
          assertEquals(Errors.NONE, resp6.error());
  
          assertEquals(EMPTY_PART_LIST, sharePartitionManager.cachedTopicIdPartitionsInShareSession(groupId, memberId2));
      }
  
      @Test
      public void testSharePartitionKey() {
          SharePartitionKey sharePartitionKey1 = new SharePartitionKey("mock-group-1",
                  new TopicIdPartition(new Uuid(0L, 1L), new TopicPartition("test", 0)));
          SharePartitionKey sharePartitionKey2 = new SharePartitionKey("mock-group-2",
                  new TopicIdPartition(new Uuid(0L, 1L), new TopicPartition("test", 0)));
          SharePartitionKey sharePartitionKey3 = new SharePartitionKey("mock-group-1",
                  new TopicIdPartition(new Uuid(1L, 1L), new TopicPartition("test-1", 0)));
          SharePartitionKey sharePartitionKey4 = new SharePartitionKey("mock-group-1",
                  new TopicIdPartition(new Uuid(0L, 1L), new TopicPartition("test", 1)));
          SharePartitionKey sharePartitionKey5 = new SharePartitionKey("mock-group-1",
                  new TopicIdPartition(new Uuid(0L, 0L), new TopicPartition("test-2", 0)));
          SharePartitionKey sharePartitionKey1Copy = new SharePartitionKey("mock-group-1",
                  new TopicIdPartition(new Uuid(0L, 1L), new TopicPartition("test", 0)));
  
          assertEquals(sharePartitionKey1, sharePartitionKey1Copy);
          assertNotEquals(sharePartitionKey1, sharePartitionKey2);
          assertNotEquals(sharePartitionKey1, sharePartitionKey3);
          assertNotEquals(sharePartitionKey1, sharePartitionKey4);
          assertNotEquals(sharePartitionKey1, sharePartitionKey5);
          assertNotNull(sharePartitionKey1);
      }
  
      @Test
      public void testMultipleSequentialShareFetches() {
          String groupId = "grp";
          Uuid memberId1 = Uuid.randomUuid();
          Uuid fooId = Uuid.randomUuid();
          Uuid barId = Uuid.randomUuid();
          TopicIdPartition tp0 = new TopicIdPartition(fooId, new TopicPartition("foo", 0));
          TopicIdPartition tp1 = new TopicIdPartition(fooId, new TopicPartition("foo", 1));
          TopicIdPartition tp2 = new TopicIdPartition(barId, new TopicPartition("bar", 0));
          TopicIdPartition tp3 = new TopicIdPartition(barId, new TopicPartition("bar", 1));
          TopicIdPartition tp4 = new TopicIdPartition(fooId, new TopicPartition("foo", 2));
          TopicIdPartition tp5 = new TopicIdPartition(barId, new TopicPartition("bar", 2));
          TopicIdPartition tp6 = new TopicIdPartition(fooId, new TopicPartition("foo", 3));
          List<TopicIdPartition> topicIdPartitions = List.of(tp0, tp1, tp2, tp3, tp4, tp5, tp6);
  
          mockFetchOffsetForTimestamp(mockReplicaManager);
  
          Timer mockTimer = systemTimerReaper();
          DelayedOperationPurgatory<DelayedShareFetch> delayedShareFetchPurgatory = new DelayedOperationPurgatory<>(
              "TestShareFetch", mockTimer, mockReplicaManager.localBrokerId(),
              DELAYED_SHARE_FETCH_PURGATORY_PURGE_INTERVAL, false, true);
          mockReplicaManagerDelayedShareFetch(mockReplicaManager, delayedShareFetchPurgatory);
          mockTopicIdPartitionToReturnDataEqualToMinBytes(mockReplicaManager, tp0, 1);
          mockTopicIdPartitionToReturnDataEqualToMinBytes(mockReplicaManager, tp1, 1);
          mockTopicIdPartitionToReturnDataEqualToMinBytes(mockReplicaManager, tp2, 1);
          mockTopicIdPartitionToReturnDataEqualToMinBytes(mockReplicaManager, tp3, 1);
          mockTopicIdPartitionToReturnDataEqualToMinBytes(mockReplicaManager, tp4, 1);
          mockTopicIdPartitionToReturnDataEqualToMinBytes(mockReplicaManager, tp5, 1);
          mockTopicIdPartitionToReturnDataEqualToMinBytes(mockReplicaManager, tp6, 1);
  
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withReplicaManager(mockReplicaManager)
              .withTimer(mockTimer)
              .withBrokerTopicStats(brokerTopicStats)
              .build();
  
          doAnswer(invocation -> buildLogReadResult(topicIdPartitions)).when(mockReplicaManager).readFromLog(any(), any(), any(ReplicaQuota.class), anyBoolean());
  
          CompletableFuture<Map<TopicIdPartition, PartitionData>> future = sharePartitionManager.fetchMessages(
              groupId, memberId1.toString(), FETCH_PARAMS, 1, MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
          assertTrue(future.isDone());
          Mockito.verify(mockReplicaManager, times(1)).readFromLog(
              any(), any(), any(ReplicaQuota.class), anyBoolean());
  
          future = sharePartitionManager.fetchMessages(groupId, memberId1.toString(), FETCH_PARAMS, 3,
              MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
          assertTrue(future.isDone());
          Mockito.verify(mockReplicaManager, times(2)).readFromLog(
              any(), any(), any(ReplicaQuota.class), anyBoolean());
  
          future = sharePartitionManager.fetchMessages(groupId, memberId1.toString(), FETCH_PARAMS, 10,
              MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
          assertTrue(future.isDone());
          Mockito.verify(mockReplicaManager, times(3)).readFromLog(
              any(), any(), any(ReplicaQuota.class), anyBoolean());
  
          // Should have 6 total fetches, 3 fetches for topic foo (though 4 partitions but 3 fetches) and 3
          // fetches for topic bar (though 3 partitions but 3 fetches).
          validateBrokerTopicStatsMetrics(
              brokerTopicStats,
              new TopicMetrics(6, 0, 0, 0),
              Map.of("foo", new TopicMetrics(3, 0, 0, 0), "bar", new TopicMetrics(3, 0, 0, 0))
          );
      }
  
      @Test
      public void testReplicaManagerFetchShouldNotProceed() {
          String groupId = "grp";
          Uuid memberId = Uuid.randomUuid();
          Uuid fooId = Uuid.randomUuid();
          TopicIdPartition tp0 = new TopicIdPartition(fooId, new TopicPartition("foo", 0));
          List<TopicIdPartition> topicIdPartitions = List.of(tp0);
  
          SharePartition sp0 = mock(SharePartition.class);
          when(sp0.maybeAcquireFetchLock(any())).thenReturn(true);
          when(sp0.canAcquireRecords()).thenReturn(false);
          when(sp0.maybeInitialize()).thenReturn(CompletableFuture.completedFuture(null));
          SharePartitionCache partitionCache = new SharePartitionCache();
          partitionCache.put(new SharePartitionKey(groupId, tp0), sp0);
  
          Timer mockTimer = systemTimerReaper();
          DelayedOperationPurgatory<DelayedShareFetch> delayedShareFetchPurgatory = new DelayedOperationPurgatory<>(
                  "TestShareFetch", mockTimer, mockReplicaManager.localBrokerId(),
                  DELAYED_SHARE_FETCH_PURGATORY_PURGE_INTERVAL, false, true);
          mockReplicaManagerDelayedShareFetch(mockReplicaManager, delayedShareFetchPurgatory);
  
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withPartitionCache(partitionCache)
              .withReplicaManager(mockReplicaManager)
              .withTimer(mockTimer)
              .withBrokerTopicStats(brokerTopicStats)
              .build();
  
          CompletableFuture<Map<TopicIdPartition, ShareFetchResponseData.PartitionData>> future =
              sharePartitionManager.fetchMessages(groupId, memberId.toString(), FETCH_PARAMS, 0,
                  MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
          Mockito.verify(mockReplicaManager, times(0)).readFromLog(
              any(), any(), any(ReplicaQuota.class), anyBoolean());
          Map<TopicIdPartition, ShareFetchResponseData.PartitionData> result = future.join();
          assertEquals(0, result.size());
          // Should have 1 fetch recorded and no failed as the fetch did complete without error.
          validateBrokerTopicStatsMetrics(
              brokerTopicStats,
              new TopicMetrics(1, 0, 0, 0),
              Map.of("foo", new TopicMetrics(1, 0, 0, 0))
          );
      }
  
      @Test
      public void testReplicaManagerFetchShouldProceed() {
          String groupId = "grp";
          Uuid memberId = Uuid.randomUuid();
          Uuid fooId = Uuid.randomUuid();
          TopicIdPartition tp0 = new TopicIdPartition(fooId, new TopicPartition("foo", 0));
          List<TopicIdPartition> topicIdPartitions = List.of(tp0);
  
          mockFetchOffsetForTimestamp(mockReplicaManager);
  
          Timer mockTimer = systemTimerReaper();
          DelayedOperationPurgatory<DelayedShareFetch> delayedShareFetchPurgatory = new DelayedOperationPurgatory<>(
              "TestShareFetch", mockTimer, mockReplicaManager.localBrokerId(),
              DELAYED_SHARE_FETCH_PURGATORY_PURGE_INTERVAL, false, true);
          mockReplicaManagerDelayedShareFetch(mockReplicaManager, delayedShareFetchPurgatory);
          mockTopicIdPartitionToReturnDataEqualToMinBytes(mockReplicaManager, tp0, 1);
  
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withReplicaManager(mockReplicaManager)
              .withTimer(mockTimer)
              .withBrokerTopicStats(brokerTopicStats)
              .build();
  
          doAnswer(invocation -> buildLogReadResult(topicIdPartitions)).when(mockReplicaManager).readFromLog(any(), any(), any(ReplicaQuota.class), anyBoolean());
  
          sharePartitionManager.fetchMessages(groupId, memberId.toString(), FETCH_PARAMS, 0,
              MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
          // Since the nextFetchOffset does not point to endOffset + 1, i.e. some of the records in the cachedState are AVAILABLE,
          // even though the maxInFlightMessages limit is exceeded, replicaManager.readFromLog should be called
          Mockito.verify(mockReplicaManager, times(1)).readFromLog(
              any(), any(), any(ReplicaQuota.class), anyBoolean());
          // Should have 1 fetch recorded.
          assertEquals(1, brokerTopicStats.allTopicsStats().totalShareFetchRequestRate().count());
          assertEquals(1, brokerTopicStats.numTopics());
          assertEquals(1, brokerTopicStats.topicStats(tp0.topic()).totalShareFetchRequestRate().count());
      }
  
      @Test
      public void testCloseSharePartitionManager() throws Exception {
          Timer timer = Mockito.mock(SystemTimerReaper.class);
          ShareGroupMetrics shareGroupMetrics = Mockito.mock(ShareGroupMetrics.class);
          SharePartitionManager sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withTimer(timer)
              .withShareGroupMetrics(shareGroupMetrics)
              .build();
  
          // Verify that 0 calls are made to timer.close() and shareGroupMetrics.close().
          Mockito.verify(timer, times(0)).close();
          Mockito.verify(shareGroupMetrics, times(0)).close();
          // Closing the sharePartitionManager closes timer object in sharePartitionManager.
          sharePartitionManager.close();
          // Verify that the timer object in sharePartitionManager is closed by checking the calls to timer.close() and shareGroupMetrics.close().
          Mockito.verify(timer, times(1)).close();
          Mockito.verify(shareGroupMetrics, times(1)).close();
      }
  
      @Test
      public void testReleaseSessionSuccess() {
          String groupId = "grp";
          Uuid memberId = Uuid.randomUuid();
  
          TopicIdPartition tp1 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 0));
          TopicIdPartition tp2 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("bar", 2));
          TopicIdPartition tp3 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("baz", 4));
  
          SharePartition sp1 = mock(SharePartition.class);
          SharePartition sp2 = mock(SharePartition.class);
          when(sp1.releaseAcquiredRecords(ArgumentMatchers.eq(memberId.toString()))).thenReturn(CompletableFuture.completedFuture(null));
          when(sp2.releaseAcquiredRecords(ArgumentMatchers.eq(memberId.toString()))).thenReturn(FutureUtils.failedFuture(
                  new InvalidRecordStateException("Unable to release acquired records for the batch")
          ));
  
          ShareSessionCache cache = mock(ShareSessionCache.class);
          ShareSession shareSession = mock(ShareSession.class);
          when(cache.get(new ShareSessionKey(groupId, memberId))).thenReturn(shareSession);
          when(cache.remove(new ShareSessionKey(groupId, memberId))).thenReturn(shareSession);
  
          ImplicitLinkedHashCollection<CachedSharePartition> partitionMap = new ImplicitLinkedHashCollection<>(3);
          partitionMap.add(new CachedSharePartition(tp1));
          partitionMap.add(new CachedSharePartition(tp2));
          partitionMap.add(new CachedSharePartition(tp3));
          when(shareSession.partitionMap()).thenReturn(partitionMap);
  
          SharePartitionCache partitionCache = new SharePartitionCache();
          partitionCache.put(new SharePartitionKey(groupId, tp1), sp1);
          partitionCache.put(new SharePartitionKey(groupId, tp2), sp2);
  
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withCache(cache)
              .withPartitionCache(partitionCache)
              .withBrokerTopicStats(brokerTopicStats)
              .build();
  
          CompletableFuture<Map<TopicIdPartition, ShareAcknowledgeResponseData.PartitionData>> resultFuture =
                  sharePartitionManager.releaseSession(groupId, memberId.toString());
          Map<TopicIdPartition, ShareAcknowledgeResponseData.PartitionData> result = resultFuture.join();
          assertEquals(3, result.size());
          assertTrue(result.containsKey(tp1));
          assertTrue(result.containsKey(tp2));
          assertTrue(result.containsKey(tp3));
          assertEquals(0, result.get(tp1).partitionIndex());
          assertEquals(Errors.NONE.code(), result.get(tp1).errorCode());
          assertEquals(2, result.get(tp2).partitionIndex());
          assertEquals(Errors.INVALID_RECORD_STATE.code(), result.get(tp2).errorCode());
          assertEquals("Unable to release acquired records for the batch", result.get(tp2).errorMessage());
          // tp3 was not a part of partitionCache.
          assertEquals(4, result.get(tp3).partitionIndex());
          assertEquals(Errors.UNKNOWN_TOPIC_OR_PARTITION.code(), result.get(tp3).errorCode());
          assertEquals(Errors.UNKNOWN_TOPIC_OR_PARTITION.message(), result.get(tp3).errorMessage());
          // Shouldn't have any metrics for fetch and acknowledge.
          validateBrokerTopicStatsMetrics(
              brokerTopicStats,
              new TopicMetrics(0, 0, 0, 0),
              Map.of()
          );
      }
  
      @Test
      public void testReleaseSessionWithIncorrectGroupId() {
          String groupId = "grp";
          Uuid memberId = Uuid.randomUuid();
  
          TopicIdPartition tp1 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 0));
  
          ShareSessionCache cache = mock(ShareSessionCache.class);
          ShareSession shareSession = mock(ShareSession.class);
          when(cache.get(new ShareSessionKey(groupId, memberId))).thenReturn(shareSession);
  
          ImplicitLinkedHashCollection<CachedSharePartition> partitionMap = new ImplicitLinkedHashCollection<>(3);
          partitionMap.add(new CachedSharePartition(tp1));
          when(shareSession.partitionMap()).thenReturn(partitionMap);
  
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withCache(cache)
              .build();
  
          // Calling releaseSession with incorrect groupId.
          CompletableFuture<Map<TopicIdPartition, ShareAcknowledgeResponseData.PartitionData>> resultFuture =
                  sharePartitionManager.releaseSession("grp-2", memberId.toString());
          assertTrue(resultFuture.isDone());
          assertTrue(resultFuture.isCompletedExceptionally());
          Throwable exception = assertThrows(ExecutionException.class, resultFuture::get);
          assertInstanceOf(ShareSessionNotFoundException.class, exception.getCause());
      }
  
      @Test
      public void testReleaseSessionWithIncorrectMemberId() {
          String groupId = "grp";
          Uuid memberId = Uuid.randomUuid();
  
          TopicIdPartition tp1 = new TopicIdPartition(memberId, new TopicPartition("foo", 0));
  
          ShareSessionCache cache = mock(ShareSessionCache.class);
          ShareSession shareSession = mock(ShareSession.class);
          // Member with random Uuid so that it does not match the memberId.
          when(cache.get(new ShareSessionKey(groupId, Uuid.randomUuid()))).thenReturn(shareSession);
  
          ImplicitLinkedHashCollection<CachedSharePartition> partitionMap = new ImplicitLinkedHashCollection<>(3);
          partitionMap.add(new CachedSharePartition(tp1));
          when(shareSession.partitionMap()).thenReturn(partitionMap);
  
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withCache(cache)
              .build();
  
          CompletableFuture<Map<TopicIdPartition, ShareAcknowledgeResponseData.PartitionData>> resultFuture =
                  sharePartitionManager.releaseSession(groupId, memberId.toString());
          assertTrue(resultFuture.isDone());
          assertTrue(resultFuture.isCompletedExceptionally());
          Throwable exception = assertThrows(ExecutionException.class, resultFuture::get);
          assertInstanceOf(ShareSessionNotFoundException.class, exception.getCause());
      }
  
      @Test
      public void testReleaseSessionWithEmptyTopicPartitions() {
          String groupId = "grp";
          Uuid memberId = Uuid.randomUuid();
  
          ShareSessionCache cache = mock(ShareSessionCache.class);
          ShareSession shareSession = mock(ShareSession.class);
          when(cache.get(new ShareSessionKey(groupId, memberId))).thenReturn(shareSession);
          when(cache.remove(new ShareSessionKey(groupId, memberId))).thenReturn(shareSession);
          when(shareSession.partitionMap()).thenReturn(new ImplicitLinkedHashCollection<>());
  
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withCache(cache)
              .build();
  
          // Empty list of TopicIdPartitions to releaseSession. This should return an empty map.
          CompletableFuture<Map<TopicIdPartition, ShareAcknowledgeResponseData.PartitionData>> resultFuture =
                  sharePartitionManager.releaseSession(groupId, memberId.toString());
          Map<TopicIdPartition, ShareAcknowledgeResponseData.PartitionData> result = resultFuture.join();
          assertEquals(0, result.size());
      }
  
      @Test
      public void testReleaseSessionWithNullShareSession() {
          String groupId = "grp";
          Uuid memberId = Uuid.randomUuid();
  
          ShareSessionCache cache = mock(ShareSessionCache.class);
          // Null share session in get response so empty topic partitions should be returned.
          when(cache.get(new ShareSessionKey(groupId, memberId))).thenReturn(null);
          // Make the response not null for remove so can further check for the return value from topic partitions.
          when(cache.remove(new ShareSessionKey(groupId, memberId))).thenReturn(mock(ShareSession.class));
  
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withCache(cache)
              .build();
  
          CompletableFuture<Map<TopicIdPartition, ShareAcknowledgeResponseData.PartitionData>> resultFuture =
              sharePartitionManager.releaseSession(groupId, memberId.toString());
          Map<TopicIdPartition, ShareAcknowledgeResponseData.PartitionData> result = resultFuture.join();
          assertEquals(0, result.size());
      }
  
      @Test
      public void testAcknowledgeSinglePartition() {
          String groupId = "grp";
          String memberId = Uuid.randomUuid().toString();
  
          TopicIdPartition tp = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 0));
          SharePartition sp = mock(SharePartition.class);
  
          when(sp.acknowledge(ArgumentMatchers.eq(memberId), any())).thenReturn(CompletableFuture.completedFuture(null));
  
          SharePartitionCache partitionCache = new SharePartitionCache();
          partitionCache.put(new SharePartitionKey(groupId, tp), sp);
  
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withPartitionCache(partitionCache)
              .withBrokerTopicStats(brokerTopicStats)
              .build();
  
          Map<TopicIdPartition, List<ShareAcknowledgementBatch>> acknowledgeTopics = new HashMap<>();
          acknowledgeTopics.put(tp, List.of(
              new ShareAcknowledgementBatch(12, 20, List.of((byte) 1)),
              new ShareAcknowledgementBatch(24, 56, List.of((byte) 1))
          ));
          CompletableFuture<Map<TopicIdPartition, ShareAcknowledgeResponseData.PartitionData>> resultFuture =
                  sharePartitionManager.acknowledge(memberId, groupId, acknowledgeTopics);
          Map<TopicIdPartition, ShareAcknowledgeResponseData.PartitionData> result = resultFuture.join();
          assertEquals(1, result.size());
          assertTrue(result.containsKey(tp));
          assertEquals(0, result.get(tp).partitionIndex());
          assertEquals(Errors.NONE.code(), result.get(tp).errorCode());
  
          validateBrokerTopicStatsMetrics(
              brokerTopicStats,
              new TopicMetrics(0, 0, 1, 0),
              Map.of("foo", new TopicMetrics(0, 0, 1, 0))
          );
      }
  
      @Test
      public void testAcknowledgeMultiplePartition() throws Exception {
          String groupId = "grp";
          String memberId = Uuid.randomUuid().toString();
  
          TopicIdPartition tp1 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo1", 0));
          TopicIdPartition tp2 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo2", 0));
          TopicIdPartition tp3 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo3", 0));
  
          SharePartition sp1 = mock(SharePartition.class);
          SharePartition sp2 = mock(SharePartition.class);
          SharePartition sp3 = mock(SharePartition.class);
  
          when(sp1.acknowledge(ArgumentMatchers.eq(memberId), any())).thenReturn(CompletableFuture.completedFuture(null));
          when(sp2.acknowledge(ArgumentMatchers.eq(memberId), any())).thenReturn(CompletableFuture.completedFuture(null));
          when(sp3.acknowledge(ArgumentMatchers.eq(memberId), any())).thenReturn(CompletableFuture.completedFuture(null));
  
          SharePartitionCache partitionCache = new SharePartitionCache();
          partitionCache.put(new SharePartitionKey(groupId, tp1), sp1);
          partitionCache.put(new SharePartitionKey(groupId, tp2), sp2);
          partitionCache.put(new SharePartitionKey(groupId, tp3), sp3);
  
          ShareGroupMetrics shareGroupMetrics = new ShareGroupMetrics(time);
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withPartitionCache(partitionCache)
              .withShareGroupMetrics(shareGroupMetrics)
              .withBrokerTopicStats(brokerTopicStats)
              .build();
  
          Map<TopicIdPartition, List<ShareAcknowledgementBatch>> acknowledgeTopics = new HashMap<>();
          acknowledgeTopics.put(tp1, List.of(
                  new ShareAcknowledgementBatch(12, 20, List.of((byte) 1)),
                  new ShareAcknowledgementBatch(24, 56, List.of((byte) 1))
          ));
          acknowledgeTopics.put(tp2, List.of(
                  new ShareAcknowledgementBatch(15, 26, List.of((byte) 2)),
                  new ShareAcknowledgementBatch(34, 56, List.of((byte) 2))
          ));
          acknowledgeTopics.put(tp3, List.of(
                  new ShareAcknowledgementBatch(4, 15, List.of((byte) 3)),
                  new ShareAcknowledgementBatch(16, 21, List.of((byte) 3))
          ));
          CompletableFuture<Map<TopicIdPartition, ShareAcknowledgeResponseData.PartitionData>> resultFuture =
                  sharePartitionManager.acknowledge(memberId, groupId, acknowledgeTopics);
          Map<TopicIdPartition, ShareAcknowledgeResponseData.PartitionData> result = resultFuture.join();
          assertEquals(3, result.size());
          assertTrue(result.containsKey(tp1));
          assertTrue(result.containsKey(tp2));
          assertTrue(result.containsKey(tp3));
          assertEquals(0, result.get(tp1).partitionIndex());
          assertEquals(Errors.NONE.code(), result.get(tp1).errorCode());
          assertEquals(0, result.get(tp2).partitionIndex());
          assertEquals(Errors.NONE.code(), result.get(tp2).errorCode());
          assertEquals(0, result.get(tp3).partitionIndex());
          assertEquals(Errors.NONE.code(), result.get(tp3).errorCode());
  
          assertEquals(42, shareGroupMetrics.recordAcknowledgementMeter(AcknowledgeType.ACCEPT.id).count());
          assertEquals(35, shareGroupMetrics.recordAcknowledgementMeter(AcknowledgeType.RELEASE.id).count());
          assertEquals(18, shareGroupMetrics.recordAcknowledgementMeter(AcknowledgeType.REJECT.id).count());
          assertTrue(shareGroupMetrics.recordAcknowledgementMeter(AcknowledgeType.ACCEPT.id).meanRate() > 0);
          assertTrue(shareGroupMetrics.recordAcknowledgementMeter(AcknowledgeType.RELEASE.id).meanRate() > 0);
          assertTrue(shareGroupMetrics.recordAcknowledgementMeter(AcknowledgeType.REJECT.id).meanRate() > 0);
  
          // Should have 3 successful acknowledgement and 1 successful acknowledgement per topic.
          validateBrokerTopicStatsMetrics(
              brokerTopicStats,
              new TopicMetrics(0, 0, 3, 0),
              Map.of(tp1.topic(), new TopicMetrics(0, 0, 1, 0), tp2.topic(), new TopicMetrics(0, 0, 1, 0), tp3.topic(), new TopicMetrics(0, 0, 1, 0))
          );
          shareGroupMetrics.close();
      }
  
      @Test
      public void testAcknowledgeIndividualOffsets() throws Exception {
          String groupId = "grp";
          String memberId = Uuid.randomUuid().toString();
  
          TopicIdPartition tp1 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo1", 0));
          TopicIdPartition tp2 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo2", 0));
  
          SharePartition sp1 = mock(SharePartition.class);
          SharePartition sp2 = mock(SharePartition.class);
  
          List<ShareAcknowledgementBatch> ack1 = List.of(
              new ShareAcknowledgementBatch(12, 12, List.of((byte) 1)));
          List<ShareAcknowledgementBatch> ack2 = List.of(
              new ShareAcknowledgementBatch(15, 20, List.of((byte) 2, (byte) 3, (byte) 2, (byte) 2, (byte) 3, (byte) 2)));
          when(sp1.acknowledge(memberId, ack1)).thenReturn(CompletableFuture.completedFuture(null));
          when(sp2.acknowledge(memberId, ack2)).thenReturn(CompletableFuture.completedFuture(null));
  
          SharePartitionCache partitionCache = new SharePartitionCache();
          partitionCache.put(new SharePartitionKey(groupId, tp1), sp1);
          partitionCache.put(new SharePartitionKey(groupId, tp2), sp2);
  
          ShareGroupMetrics shareGroupMetrics = new ShareGroupMetrics(time);
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withPartitionCache(partitionCache)
              .withShareGroupMetrics(shareGroupMetrics)
              .withBrokerTopicStats(brokerTopicStats)
              .build();
  
          Map<TopicIdPartition, List<ShareAcknowledgementBatch>> acknowledgeTopics = Map.of(tp1, ack1, tp2, ack2);
          CompletableFuture<Map<TopicIdPartition, ShareAcknowledgeResponseData.PartitionData>> resultFuture =
              sharePartitionManager.acknowledge(memberId, groupId, acknowledgeTopics);
          Map<TopicIdPartition, ShareAcknowledgeResponseData.PartitionData> result = resultFuture.join();
          assertEquals(2, result.size());
          assertTrue(result.containsKey(tp1));
          assertTrue(result.containsKey(tp2));
          assertEquals(0, result.get(tp1).partitionIndex());
          assertEquals(Errors.NONE.code(), result.get(tp1).errorCode());
          assertEquals(0, result.get(tp2).partitionIndex());
          assertEquals(Errors.NONE.code(), result.get(tp2).errorCode());
  
          assertEquals(1, shareGroupMetrics.recordAcknowledgementMeter(AcknowledgeType.ACCEPT.id).count());
          assertEquals(4, shareGroupMetrics.recordAcknowledgementMeter(AcknowledgeType.RELEASE.id).count());
          assertEquals(2, shareGroupMetrics.recordAcknowledgementMeter(AcknowledgeType.REJECT.id).count());
          assertTrue(shareGroupMetrics.recordAcknowledgementMeter(AcknowledgeType.ACCEPT.id).meanRate() > 0);
          assertTrue(shareGroupMetrics.recordAcknowledgementMeter(AcknowledgeType.RELEASE.id).meanRate() > 0);
          assertTrue(shareGroupMetrics.recordAcknowledgementMeter(AcknowledgeType.REJECT.id).meanRate() > 0);
  
          shareGroupMetrics.close();
      }
  
      @Test
      public void testAcknowledgeIncorrectGroupId() {
          String groupId = "grp";
          String groupId2 = "grp2";
          String memberId = Uuid.randomUuid().toString();
  
          TopicIdPartition tp = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 0));
          SharePartition sp = mock(SharePartition.class);
  
          SharePartitionCache partitionCache = new SharePartitionCache();
          partitionCache.put(new SharePartitionKey(groupId, tp), sp);
          ShareGroupMetrics shareGroupMetrics = new ShareGroupMetrics(time);
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withPartitionCache(partitionCache)
              .withBrokerTopicStats(brokerTopicStats)
              .withShareGroupMetrics(shareGroupMetrics)
              .build();
  
          Map<TopicIdPartition, List<ShareAcknowledgementBatch>> acknowledgeTopics = new HashMap<>();
          acknowledgeTopics.put(tp, List.of(
              new ShareAcknowledgementBatch(12, 20, List.of((byte) 1)),
              new ShareAcknowledgementBatch(24, 56, List.of((byte) 1))
          ));
          CompletableFuture<Map<TopicIdPartition, ShareAcknowledgeResponseData.PartitionData>> resultFuture =
                  sharePartitionManager.acknowledge(memberId, groupId2, acknowledgeTopics);
          Map<TopicIdPartition, ShareAcknowledgeResponseData.PartitionData> result = resultFuture.join();
          assertEquals(1, result.size());
          assertTrue(result.containsKey(tp));
          assertEquals(0, result.get(tp).partitionIndex());
          assertEquals(Errors.UNKNOWN_TOPIC_OR_PARTITION.code(), result.get(tp).errorCode());
          assertEquals(Errors.UNKNOWN_TOPIC_OR_PARTITION.message(), result.get(tp).errorMessage());
          // No metric should be recorded as acknowledge failed.
          assertEquals(0, shareGroupMetrics.recordAcknowledgementMeter(AcknowledgeType.ACCEPT.id).count());
          assertEquals(0, shareGroupMetrics.recordAcknowledgementMeter(AcknowledgeType.RELEASE.id).count());
          assertEquals(0, shareGroupMetrics.recordAcknowledgementMeter(AcknowledgeType.REJECT.id).count());
          // Should have 1 acknowledge recorded and 1 failed.
          validateBrokerTopicStatsMetrics(
              brokerTopicStats,
              new TopicMetrics(0, 0, 1, 1),
              Map.of(tp.topic(), new TopicMetrics(0, 0, 1, 1))
          );
      }
  
      @Test
      public void testAcknowledgeIncorrectMemberId() {
          String groupId = "grp";
          String memberId = Uuid.randomUuid().toString();
  
          TopicIdPartition tp = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 0));
          SharePartition sp = mock(SharePartition.class);
          when(sp.acknowledge(ArgumentMatchers.eq(memberId), any())).thenReturn(FutureUtils.failedFuture(
                  new InvalidRequestException("Member is not the owner of batch record")
          ));
          SharePartitionCache partitionCache = new SharePartitionCache();
          partitionCache.put(new SharePartitionKey(groupId, tp), sp);
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withPartitionCache(partitionCache)
              .withBrokerTopicStats(brokerTopicStats)
              .build();
  
          Map<TopicIdPartition, List<ShareAcknowledgementBatch>> acknowledgeTopics = new HashMap<>();
          acknowledgeTopics.put(tp, List.of(
              new ShareAcknowledgementBatch(12, 20, List.of((byte) 1)),
              new ShareAcknowledgementBatch(24, 56, List.of((byte) 1))
          ));
  
          CompletableFuture<Map<TopicIdPartition, ShareAcknowledgeResponseData.PartitionData>> resultFuture =
              sharePartitionManager.acknowledge(memberId, groupId, acknowledgeTopics);
          Map<TopicIdPartition, ShareAcknowledgeResponseData.PartitionData> result = resultFuture.join();
          assertEquals(1, result.size());
          assertTrue(result.containsKey(tp));
          assertEquals(0, result.get(tp).partitionIndex());
          assertEquals(Errors.INVALID_REQUEST.code(), result.get(tp).errorCode());
          assertEquals("Member is not the owner of batch record", result.get(tp).errorMessage());
          // Should have 1 acknowledge recorded and 1 failed.
          validateBrokerTopicStatsMetrics(
              brokerTopicStats,
              new TopicMetrics(0, 0, 1, 1),
              Map.of(tp.topic(), new TopicMetrics(0, 0, 1, 1))
          );
      }
  
      @Test
      public void testAcknowledgeEmptyPartitionCacheMap() {
          String groupId = "grp";
          String memberId = Uuid.randomUuid().toString();
  
          TopicIdPartition tp = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo4", 3));
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withBrokerTopicStats(brokerTopicStats)
              .build();
  
          Map<TopicIdPartition, List<ShareAcknowledgementBatch>> acknowledgeTopics = new HashMap<>();
          acknowledgeTopics.put(tp, List.of(
                  new ShareAcknowledgementBatch(78, 90, List.of((byte) 2)),
                  new ShareAcknowledgementBatch(94, 99, List.of((byte) 2))
          ));
          CompletableFuture<Map<TopicIdPartition, ShareAcknowledgeResponseData.PartitionData>> resultFuture =
                  sharePartitionManager.acknowledge(memberId, groupId, acknowledgeTopics);
          Map<TopicIdPartition, ShareAcknowledgeResponseData.PartitionData> result = resultFuture.join();
          assertEquals(1, result.size());
          assertTrue(result.containsKey(tp));
          assertEquals(3, result.get(tp).partitionIndex());
          assertEquals(Errors.UNKNOWN_TOPIC_OR_PARTITION.code(), result.get(tp).errorCode());
          assertEquals(Errors.UNKNOWN_TOPIC_OR_PARTITION.message(), result.get(tp).errorMessage());
          // Should have 1 acknowledge recorded and 1 failed.
          validateBrokerTopicStatsMetrics(
              brokerTopicStats,
              new TopicMetrics(0, 0, 1, 1),
              Map.of(tp.topic(), new TopicMetrics(0, 0, 1, 1))
          );
      }
  
      @Test
      public void testAcknowledgeCompletesDelayedShareFetchRequest() {
          String groupId = "grp";
          String memberId = Uuid.randomUuid().toString();
  
          TopicIdPartition tp1 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo1", 0));
          TopicIdPartition tp2 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo2", 0));
  
          List<TopicIdPartition> topicIdPartitions = List.of(tp1, tp2);
  
          SharePartition sp1 = mock(SharePartition.class);
          SharePartition sp2 = mock(SharePartition.class);
          // mocked share partitions sp1 and sp2 can be acquired once there is an acknowledgement for it.
          doAnswer(invocation -> {
              when(sp1.canAcquireRecords()).thenReturn(true);
              return CompletableFuture.completedFuture(Optional.empty());
          }).when(sp1).acknowledge(ArgumentMatchers.eq(memberId), any());
          doAnswer(invocation -> {
              when(sp2.canAcquireRecords()).thenReturn(true);
              return CompletableFuture.completedFuture(Optional.empty());
          }).when(sp2).acknowledge(ArgumentMatchers.eq(memberId), any());
  
          SharePartitionCache partitionCache = new SharePartitionCache();
          partitionCache.put(new SharePartitionKey(groupId, tp1), sp1);
          partitionCache.put(new SharePartitionKey(groupId, tp2), sp2);
  
          ShareFetch shareFetch = new ShareFetch(
              FETCH_PARAMS,
              groupId,
              Uuid.randomUuid().toString(),
              new CompletableFuture<>(),
              topicIdPartitions,
              BATCH_SIZE,
              100,
              brokerTopicStats);
  
          Timer mockTimer = systemTimerReaper();
          DelayedOperationPurgatory<DelayedShareFetch> delayedShareFetchPurgatory = new DelayedOperationPurgatory<>(
                  "TestShareFetch", mockTimer, mockReplicaManager.localBrokerId(),
                  DELAYED_SHARE_FETCH_PURGATORY_PURGE_INTERVAL, false, true);
          mockReplicaManagerDelayedShareFetch(mockReplicaManager, delayedShareFetchPurgatory);
          when(sp1.fetchOffsetMetadata(anyLong())).thenReturn(Optional.of(new LogOffsetMetadata(0, 1, 0)));
          mockTopicIdPartitionToReturnDataEqualToMinBytes(mockReplicaManager, tp1, 2);
  
          // Initially you cannot acquire records for both sp1 and sp2.
          when(sp1.maybeAcquireFetchLock(any())).thenReturn(true);
          when(sp1.canAcquireRecords()).thenReturn(false);
          when(sp2.maybeAcquireFetchLock(any())).thenReturn(true);
          when(sp2.canAcquireRecords()).thenReturn(false);
          when(sp1.acquire(anyString(), anyInt(), anyInt(), anyLong(), any(), any())).thenReturn(ShareAcquiredRecords.empty());
          when(sp2.acquire(anyString(), anyInt(), anyInt(), anyLong(), any(), any())).thenReturn(ShareAcquiredRecords.empty());
  
          List<DelayedOperationKey> delayedShareFetchWatchKeys = new ArrayList<>();
          topicIdPartitions.forEach(topicIdPartition -> delayedShareFetchWatchKeys.add(new DelayedShareFetchGroupKey(groupId, topicIdPartition.topicId(), topicIdPartition.partition())));
  
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withPartitionCache(partitionCache)
              .withReplicaManager(mockReplicaManager)
              .withTimer(mockTimer)
              .withBrokerTopicStats(brokerTopicStats)
              .build();
  
          LinkedHashMap<TopicIdPartition, SharePartition> sharePartitions = new LinkedHashMap<>();
          sharePartitions.put(tp1, sp1);
          sharePartitions.put(tp2, sp2);
  
          DelayedShareFetch delayedShareFetch = DelayedShareFetchTest.DelayedShareFetchBuilder.builder()
              .withShareFetchData(shareFetch)
              .withReplicaManager(mockReplicaManager)
              .withSharePartitions(sharePartitions)
              .withPartitionMaxBytesStrategy(PartitionMaxBytesStrategy.type(PartitionMaxBytesStrategy.StrategyType.UNIFORM))
              .build();
  
          delayedShareFetchPurgatory.tryCompleteElseWatch(delayedShareFetch, delayedShareFetchWatchKeys);
  
          // Since acquisition lock for sp1 and sp2 cannot be acquired, we should have 2 watched keys.
          assertEquals(2, delayedShareFetchPurgatory.watched());
  
          doAnswer(invocation -> buildLogReadResult(List.of(tp1))).when(mockReplicaManager).readFromLog(any(), any(), any(ReplicaQuota.class), anyBoolean());
  
          Map<TopicIdPartition, List<ShareAcknowledgementBatch>> acknowledgeTopics = new HashMap<>();
          acknowledgeTopics.put(tp1, List.of(
                  new ShareAcknowledgementBatch(12, 20, List.of((byte) 1)),
                  new ShareAcknowledgementBatch(24, 56, List.of((byte) 1))
          ));
  
          assertEquals(2, delayedShareFetchPurgatory.watched());
          // Acknowledgement request for sp1.
          sharePartitionManager.acknowledge(memberId, groupId, acknowledgeTopics);
  
          // Since sp1 is acknowledged, the delayedShareFetchPurgatory should have 1 watched key corresponding to sp2.
          assertEquals(1, delayedShareFetchPurgatory.watched());
  
          Mockito.verify(sp1, times(1)).nextFetchOffset();
          Mockito.verify(sp2, times(0)).nextFetchOffset();
          assertTrue(delayedShareFetch.lock().tryLock());
          delayedShareFetch.lock().unlock();
          // Should have 1 acknowledge recorded as other topic is acknowledgement request is not sent.
          assertEquals(1, brokerTopicStats.allTopicsStats().totalShareAcknowledgementRequestRate().count());
          assertEquals(1, brokerTopicStats.numTopics());
          assertEquals(1, brokerTopicStats.topicStats(tp1.topic()).totalShareAcknowledgementRequestRate().count());
      }
  
      @Test
      public void testAcknowledgeDoesNotCompleteDelayedShareFetchRequest() {
          String groupId = "grp";
          String memberId = Uuid.randomUuid().toString();
  
          TopicIdPartition tp1 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo1", 0));
          TopicIdPartition tp2 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo2", 0));
          TopicIdPartition tp3 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo3", 0));
  
          List<TopicIdPartition> topicIdPartitions = List.of(tp1, tp2);
  
          SharePartition sp1 = mock(SharePartition.class);
          SharePartition sp2 = mock(SharePartition.class);
          SharePartition sp3 = mock(SharePartition.class);
  
          // mocked share partitions sp1, sp2 and sp3 can be acquired once there is an acknowledgement for it.
          doAnswer(invocation -> {
              when(sp1.canAcquireRecords()).thenReturn(true);
              return CompletableFuture.completedFuture(Optional.empty());
          }).when(sp1).acknowledge(ArgumentMatchers.eq(memberId), any());
          doAnswer(invocation -> {
              when(sp2.canAcquireRecords()).thenReturn(true);
              return CompletableFuture.completedFuture(Optional.empty());
          }).when(sp2).acknowledge(ArgumentMatchers.eq(memberId), any());
          doAnswer(invocation -> {
              when(sp3.canAcquireRecords()).thenReturn(true);
              return CompletableFuture.completedFuture(Optional.empty());
          }).when(sp3).acknowledge(ArgumentMatchers.eq(memberId), any());
  
          SharePartitionCache partitionCache = new SharePartitionCache();
          partitionCache.put(new SharePartitionKey(groupId, tp1), sp1);
          partitionCache.put(new SharePartitionKey(groupId, tp2), sp2);
          partitionCache.put(new SharePartitionKey(groupId, tp3), sp3);
  
          ShareFetch shareFetch = new ShareFetch(
              FETCH_PARAMS,
              groupId,
              Uuid.randomUuid().toString(),
              new CompletableFuture<>(),
              topicIdPartitions,
              BATCH_SIZE,
              100,
              brokerTopicStats);
  
          Timer mockTimer = systemTimerReaper();
          DelayedOperationPurgatory<DelayedShareFetch> delayedShareFetchPurgatory = new DelayedOperationPurgatory<>(
                  "TestShareFetch", mockTimer, mockReplicaManager.localBrokerId(),
                  DELAYED_SHARE_FETCH_PURGATORY_PURGE_INTERVAL, false, true);
          mockReplicaManagerDelayedShareFetch(mockReplicaManager, delayedShareFetchPurgatory);
  
          // Initially you cannot acquire records for both all 3 share partitions.
          when(sp1.maybeAcquireFetchLock(any())).thenReturn(true);
          when(sp1.canAcquireRecords()).thenReturn(false);
          when(sp2.maybeAcquireFetchLock(any())).thenReturn(true);
          when(sp2.canAcquireRecords()).thenReturn(false);
          when(sp3.maybeAcquireFetchLock(any())).thenReturn(true);
          when(sp3.canAcquireRecords()).thenReturn(false);
  
          List<DelayedOperationKey> delayedShareFetchWatchKeys = new ArrayList<>();
          topicIdPartitions.forEach(topicIdPartition -> delayedShareFetchWatchKeys.add(new DelayedShareFetchGroupKey(groupId, topicIdPartition.topicId(), topicIdPartition.partition())));
  
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withPartitionCache(partitionCache)
              .withReplicaManager(mockReplicaManager)
              .withTimer(mockTimer)
              .withBrokerTopicStats(brokerTopicStats)
              .build();
  
          LinkedHashMap<TopicIdPartition, SharePartition> sharePartitions = new LinkedHashMap<>();
          sharePartitions.put(tp1, sp1);
          sharePartitions.put(tp2, sp2);
          sharePartitions.put(tp3, sp3);
  
          DelayedShareFetch delayedShareFetch = DelayedShareFetchTest.DelayedShareFetchBuilder.builder()
              .withShareFetchData(shareFetch)
              .withReplicaManager(mockReplicaManager)
              .withSharePartitions(sharePartitions)
              .build();
  
          delayedShareFetchPurgatory.tryCompleteElseWatch(delayedShareFetch, delayedShareFetchWatchKeys);
  
          // Since acquisition lock for sp1 and sp2 cannot be acquired, we should have 2 watched keys.
          assertEquals(2, delayedShareFetchPurgatory.watched());
  
          Map<TopicIdPartition, List<ShareAcknowledgementBatch>> acknowledgeTopics = new HashMap<>();
          acknowledgeTopics.put(tp3, List.of(
                  new ShareAcknowledgementBatch(12, 20, List.of((byte) 1)),
                  new ShareAcknowledgementBatch(24, 56, List.of((byte) 1))
          ));
  
          // Acknowledgement request for sp3.
          sharePartitionManager.acknowledge(memberId, groupId, acknowledgeTopics);
  
          // Since neither sp1 and sp2 have been acknowledged, the delayedShareFetchPurgatory should have 2 watched keys.
          assertEquals(2, delayedShareFetchPurgatory.watched());
  
          Mockito.verify(sp1, times(0)).nextFetchOffset();
          Mockito.verify(sp2, times(0)).nextFetchOffset();
          assertTrue(delayedShareFetch.lock().tryLock());
          delayedShareFetch.lock().unlock();
          // Should have 1 acknowledge recorded as other 2 topics acknowledgement request is not sent.
          assertEquals(1, brokerTopicStats.allTopicsStats().totalShareAcknowledgementRequestRate().count());
          assertEquals(1, brokerTopicStats.numTopics());
          assertEquals(1, brokerTopicStats.topicStats(tp3.topic()).totalShareAcknowledgementRequestRate().count());
      }
  
      @Test
      public void testReleaseSessionCompletesDelayedShareFetchRequest() {
          String groupId = "grp";
          String memberId = Uuid.randomUuid().toString();
  
          TopicIdPartition tp1 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo1", 0));
          TopicIdPartition tp2 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo2", 0));
          TopicIdPartition tp3 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo3", 0));
  
          List<TopicIdPartition> topicIdPartitions = List.of(tp1, tp2);
  
          SharePartition sp1 = mock(SharePartition.class);
          SharePartition sp2 = mock(SharePartition.class);
          SharePartition sp3 = mock(SharePartition.class);
  
          ShareSessionCache cache = mock(ShareSessionCache.class);
          ShareSession shareSession = mock(ShareSession.class);
          when(cache.remove(new ShareSessionKey(groupId, Uuid.fromString(memberId)))).thenReturn(shareSession);
  
          // mocked share partitions sp1 and sp2 can be acquired once there is a release acquired records on session close request for it.
          doAnswer(invocation -> {
              when(sp1.canAcquireRecords()).thenReturn(true);
              return CompletableFuture.completedFuture(Optional.empty());
          }).when(sp1).releaseAcquiredRecords(ArgumentMatchers.eq(memberId));
          doAnswer(invocation -> {
              when(sp2.canAcquireRecords()).thenReturn(true);
              return CompletableFuture.completedFuture(Optional.empty());
          }).when(sp2).releaseAcquiredRecords(ArgumentMatchers.eq(memberId));
          when(sp3.releaseAcquiredRecords(ArgumentMatchers.eq(memberId))).thenReturn(CompletableFuture.completedFuture(null));
  
          SharePartitionCache partitionCache = new SharePartitionCache();
          partitionCache.put(new SharePartitionKey(groupId, tp1), sp1);
          partitionCache.put(new SharePartitionKey(groupId, tp2), sp2);
          partitionCache.put(new SharePartitionKey(groupId, tp3), sp3);
  
          ShareFetch shareFetch = new ShareFetch(
              FETCH_PARAMS,
              groupId,
              Uuid.randomUuid().toString(),
              new CompletableFuture<>(),
              topicIdPartitions,
              BATCH_SIZE,
              100,
              brokerTopicStats);
  
          Timer mockTimer = systemTimerReaper();
          DelayedOperationPurgatory<DelayedShareFetch> delayedShareFetchPurgatory = new DelayedOperationPurgatory<>(
                  "TestShareFetch", mockTimer, mockReplicaManager.localBrokerId(),
                  DELAYED_SHARE_FETCH_PURGATORY_PURGE_INTERVAL, false, true);
          mockReplicaManagerDelayedShareFetch(mockReplicaManager, delayedShareFetchPurgatory);
          when(sp1.fetchOffsetMetadata(anyLong())).thenReturn(Optional.of(new LogOffsetMetadata(0, 1, 0)));
          mockTopicIdPartitionToReturnDataEqualToMinBytes(mockReplicaManager, tp1, 1);
  
          // Initially you cannot acquire records for both sp1 and sp2.
          when(sp1.maybeAcquireFetchLock(any())).thenReturn(true);
          when(sp1.canAcquireRecords()).thenReturn(false);
          when(sp2.maybeAcquireFetchLock(any())).thenReturn(true);
          when(sp2.canAcquireRecords()).thenReturn(false);
  
          List<DelayedOperationKey> delayedShareFetchWatchKeys = new ArrayList<>();
          topicIdPartitions.forEach(topicIdPartition -> delayedShareFetchWatchKeys.add(new DelayedShareFetchGroupKey(groupId, topicIdPartition.topicId(), topicIdPartition.partition())));
  
          sharePartitionManager = spy(SharePartitionManagerBuilder.builder()
              .withPartitionCache(partitionCache)
              .withCache(cache)
              .withReplicaManager(mockReplicaManager)
              .withTimer(mockTimer)
              .build());
  
          LinkedHashMap<TopicIdPartition, SharePartition> sharePartitions = new LinkedHashMap<>();
          sharePartitions.put(tp1, sp1);
          sharePartitions.put(tp2, sp2);
  
          DelayedShareFetch delayedShareFetch = DelayedShareFetchTest.DelayedShareFetchBuilder.builder()
              .withShareFetchData(shareFetch)
              .withReplicaManager(mockReplicaManager)
              .withSharePartitions(sharePartitions)
              .withPartitionMaxBytesStrategy(PartitionMaxBytesStrategy.type(PartitionMaxBytesStrategy.StrategyType.UNIFORM))
              .build();
  
          delayedShareFetchPurgatory.tryCompleteElseWatch(delayedShareFetch, delayedShareFetchWatchKeys);
  
          // Since acquisition lock for sp1 and sp2 cannot be acquired, we should have 2 watched keys.
          assertEquals(2, delayedShareFetchPurgatory.watched());
  
          // The share session for this share group member returns tp1 and tp3, tp1 is common in both the delayed fetch request and the share session.
          when(sharePartitionManager.cachedTopicIdPartitionsInShareSession(groupId, Uuid.fromString(memberId))).thenReturn(List.of(tp1, tp3));
  
          doAnswer(invocation -> buildLogReadResult(List.of(tp1))).when(mockReplicaManager).readFromLog(any(), any(), any(ReplicaQuota.class), anyBoolean());
          when(sp1.acquire(anyString(), anyInt(), anyInt(), anyLong(), any(), any())).thenReturn(new ShareAcquiredRecords(EMPTY_ACQUIRED_RECORDS, 0));
          // Release acquired records on session close request for tp1 and tp3.
          sharePartitionManager.releaseSession(groupId, memberId);
  
          // Since sp1's request to release acquired records on session close is completed, the delayedShareFetchPurgatory
          // should have 1 watched key corresponding to sp2.
          assertEquals(1, delayedShareFetchPurgatory.watched());
  
          Mockito.verify(sp1, times(1)).nextFetchOffset();
          Mockito.verify(sp2, times(0)).nextFetchOffset();
          assertTrue(delayedShareFetch.lock().tryLock());
          delayedShareFetch.lock().unlock();
      }
  
      @Test
      public void testReleaseSessionDoesNotCompleteDelayedShareFetchRequest() {
          String groupId = "grp";
          String memberId = Uuid.randomUuid().toString();
  
          TopicIdPartition tp1 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo1", 0));
          TopicIdPartition tp2 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo2", 0));
          TopicIdPartition tp3 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo3", 0));
  
          List<TopicIdPartition> topicIdPartitions = List.of(tp1, tp2);
  
          SharePartition sp1 = mock(SharePartition.class);
          SharePartition sp2 = mock(SharePartition.class);
          SharePartition sp3 = mock(SharePartition.class);
  
          ShareSessionCache cache = mock(ShareSessionCache.class);
          ShareSession shareSession = mock(ShareSession.class);
          when(cache.remove(new ShareSessionKey(groupId, Uuid.fromString(memberId)))).thenReturn(shareSession);
  
          // mocked share partitions sp1, sp2 and sp3 can be acquired once there is a release acquired records on session close for it.
          doAnswer(invocation -> {
              when(sp1.canAcquireRecords()).thenReturn(true);
              return CompletableFuture.completedFuture(Optional.empty());
          }).when(sp1).releaseAcquiredRecords(ArgumentMatchers.eq(memberId));
          doAnswer(invocation -> {
              when(sp2.canAcquireRecords()).thenReturn(true);
              return CompletableFuture.completedFuture(Optional.empty());
          }).when(sp2).releaseAcquiredRecords(ArgumentMatchers.eq(memberId));
          doAnswer(invocation -> {
              when(sp3.canAcquireRecords()).thenReturn(true);
              return CompletableFuture.completedFuture(Optional.empty());
          }).when(sp3).releaseAcquiredRecords(ArgumentMatchers.eq(memberId));
  
          SharePartitionCache partitionCache = new SharePartitionCache();
          partitionCache.put(new SharePartitionKey(groupId, tp1), sp1);
          partitionCache.put(new SharePartitionKey(groupId, tp2), sp2);
          partitionCache.put(new SharePartitionKey(groupId, tp3), sp3);
  
          ShareFetch shareFetch = new ShareFetch(
              FETCH_PARAMS,
              groupId,
              Uuid.randomUuid().toString(),
              new CompletableFuture<>(),
              topicIdPartitions,
              BATCH_SIZE,
              100,
              brokerTopicStats);
  
          Timer mockTimer = systemTimerReaper();
          DelayedOperationPurgatory<DelayedShareFetch> delayedShareFetchPurgatory = new DelayedOperationPurgatory<>(
                  "TestShareFetch", mockTimer, mockReplicaManager.localBrokerId(),
                  DELAYED_SHARE_FETCH_PURGATORY_PURGE_INTERVAL, false, true);
          mockReplicaManagerDelayedShareFetch(mockReplicaManager, delayedShareFetchPurgatory);
  
          // Initially you cannot acquire records for both all 3 share partitions.
          when(sp1.maybeAcquireFetchLock(any())).thenReturn(true);
          when(sp1.canAcquireRecords()).thenReturn(false);
          when(sp2.maybeAcquireFetchLock(any())).thenReturn(true);
          when(sp2.canAcquireRecords()).thenReturn(false);
          when(sp3.maybeAcquireFetchLock(any())).thenReturn(true);
          when(sp3.canAcquireRecords()).thenReturn(false);
  
          List<DelayedOperationKey> delayedShareFetchWatchKeys = new ArrayList<>();
          topicIdPartitions.forEach(topicIdPartition -> delayedShareFetchWatchKeys.add(new DelayedShareFetchGroupKey(groupId, topicIdPartition.topicId(), topicIdPartition.partition())));
  
          sharePartitionManager = spy(SharePartitionManagerBuilder.builder()
              .withPartitionCache(partitionCache)
              .withCache(cache)
              .withReplicaManager(mockReplicaManager)
              .withTimer(mockTimer)
              .build());
  
          LinkedHashMap<TopicIdPartition, SharePartition> sharePartitions = new LinkedHashMap<>();
          sharePartitions.put(tp1, sp1);
          sharePartitions.put(tp2, sp2);
          sharePartitions.put(tp3, sp3);
  
          DelayedShareFetch delayedShareFetch = DelayedShareFetchTest.DelayedShareFetchBuilder.builder()
              .withShareFetchData(shareFetch)
              .withReplicaManager(mockReplicaManager)
              .withSharePartitions(sharePartitions)
              .build();
  
          delayedShareFetchPurgatory.tryCompleteElseWatch(delayedShareFetch, delayedShareFetchWatchKeys);
  
          // Since acquisition lock for sp1 and sp2 cannot be acquired, we should have 2 watched keys.
          assertEquals(2, delayedShareFetchPurgatory.watched());
  
          // The share session for this share group member returns tp1 and tp3. No topic partition is common in
          // both the delayed fetch request and the share session.
          when(sharePartitionManager.cachedTopicIdPartitionsInShareSession(groupId, Uuid.fromString(memberId))).thenReturn(List.of(tp3));
  
          // Release acquired records on session close for sp3.
          sharePartitionManager.releaseSession(groupId, memberId);
  
          // Since neither sp1 and sp2 are a part of the release acquired records request on session close, the
          // delayedShareFetchPurgatory should have 2 watched keys.
          assertEquals(2, delayedShareFetchPurgatory.watched());
  
          Mockito.verify(sp1, times(0)).nextFetchOffset();
          Mockito.verify(sp2, times(0)).nextFetchOffset();
          assertTrue(delayedShareFetch.lock().tryLock());
          delayedShareFetch.lock().unlock();
      }
  
      @Test
      public void testPendingInitializationShouldCompleteFetchRequest() throws Exception {
          String groupId = "grp";
          Uuid memberId = Uuid.randomUuid();
          Uuid fooId = Uuid.randomUuid();
          TopicIdPartition tp0 = new TopicIdPartition(fooId, new TopicPartition("foo", 0));
          List<TopicIdPartition> topicIdPartitions = List.of(tp0);
  
          SharePartition sp0 = mock(SharePartition.class);
          SharePartitionCache partitionCache = new SharePartitionCache();
          partitionCache.put(new SharePartitionKey(groupId, tp0), sp0);
  
          // Keep the initialization future pending, so fetch request is stuck.
          CompletableFuture<Void> pendingInitializationFuture = new CompletableFuture<>();
          when(sp0.maybeInitialize()).thenReturn(pendingInitializationFuture);
          when(sp0.loadStartTimeMs()).thenReturn(10L);
  
          Timer mockTimer = systemTimerReaper();
          DelayedOperationPurgatory<DelayedShareFetch> delayedShareFetchPurgatory = new DelayedOperationPurgatory<>(
                  "TestShareFetch", mockTimer, mockReplicaManager.localBrokerId(),
                  DELAYED_SHARE_FETCH_PURGATORY_PURGE_INTERVAL, false, true);
          mockReplicaManagerDelayedShareFetch(mockReplicaManager, delayedShareFetchPurgatory);
  
          Time time = mock(Time.class);
          when(time.hiResClockMs()).thenReturn(100L);
          ShareGroupMetrics shareGroupMetrics = new ShareGroupMetrics(time);
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withPartitionCache(partitionCache)
              .withReplicaManager(mockReplicaManager)
              .withTime(time)
              .withShareGroupMetrics(shareGroupMetrics)
              .withTimer(mockTimer)
              .withBrokerTopicStats(brokerTopicStats)
              .build();
  
          CompletableFuture<Map<TopicIdPartition, ShareFetchResponseData.PartitionData>> future =
              sharePartitionManager.fetchMessages(groupId, memberId.toString(), FETCH_PARAMS, 0,
                  MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
          // Verify that the fetch request is completed.
          TestUtils.waitForCondition(
              future::isDone,
              DELAYED_SHARE_FETCH_TIMEOUT_MS,
              () -> "Processing in delayed share fetch queue never ended.");
          assertTrue(future.join().isEmpty());
          // Verify that replica manager fetch is not called.
          Mockito.verify(mockReplicaManager, times(0)).readFromLog(
              any(), any(), any(ReplicaQuota.class), anyBoolean());
          assertFalse(pendingInitializationFuture.isDone());
          assertEquals(0, shareGroupMetrics.partitionLoadTimeMs().count());
          // Complete the pending initialization future.
          pendingInitializationFuture.complete(null);
          // Verify the partition load time metrics.
          assertEquals(1, shareGroupMetrics.partitionLoadTimeMs().count());
          assertEquals(90.0, shareGroupMetrics.partitionLoadTimeMs().min());
          assertEquals(90.0, shareGroupMetrics.partitionLoadTimeMs().max());
          assertEquals(90.0, shareGroupMetrics.partitionLoadTimeMs().sum());
          // Should have 1 fetch recorded.
          validateBrokerTopicStatsMetrics(
              brokerTopicStats,
              new TopicMetrics(1, 0, 0, 0),
              Map.of(tp0.topic(), new TopicMetrics(1, 0, 0, 0))
          );
          shareGroupMetrics.close();
      }
  
      @Test
      public void testPartitionLoadTimeMetricWithMultiplePartitions() throws Exception {
          String groupId = "grp";
          TopicIdPartition tp0 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 0));
          TopicIdPartition tp1 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 1));
          List<TopicIdPartition> topicIdPartitions = List.of(tp0, tp1);
  
          SharePartition sp0 = mock(SharePartition.class);
          SharePartition sp1 = mock(SharePartition.class);
          SharePartitionCache partitionCache = new SharePartitionCache();
          partitionCache.put(new SharePartitionKey(groupId, tp0), sp0);
          partitionCache.put(new SharePartitionKey(groupId, tp1), sp1);
  
          // Keep the initialization future pending, so fetch request is stuck.
          CompletableFuture<Void> pendingInitializationFuture1 = new CompletableFuture<>();
          when(sp0.maybeInitialize()).thenReturn(pendingInitializationFuture1);
          when(sp0.loadStartTimeMs()).thenReturn(10L);
  
          CompletableFuture<Void> pendingInitializationFuture2 = new CompletableFuture<>();
          when(sp1.maybeInitialize()).thenReturn(pendingInitializationFuture2);
          when(sp1.loadStartTimeMs()).thenReturn(40L);
  
          Timer mockTimer = systemTimerReaper();
          DelayedOperationPurgatory<DelayedShareFetch> delayedShareFetchPurgatory = new DelayedOperationPurgatory<>(
              "TestShareFetch", mockTimer, mockReplicaManager.localBrokerId(),
              DELAYED_SHARE_FETCH_PURGATORY_PURGE_INTERVAL, false, true);
          mockReplicaManagerDelayedShareFetch(mockReplicaManager, delayedShareFetchPurgatory);
  
          Time time = mock(Time.class);
          when(time.hiResClockMs()).thenReturn(100L);
          ShareGroupMetrics shareGroupMetrics = new ShareGroupMetrics(time);
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withPartitionCache(partitionCache)
              .withReplicaManager(mockReplicaManager)
              .withTime(time)
              .withShareGroupMetrics(shareGroupMetrics)
              .withTimer(mockTimer)
              .withBrokerTopicStats(brokerTopicStats)
              .build();
  
          CompletableFuture<Map<TopicIdPartition, ShareFetchResponseData.PartitionData>> future =
              sharePartitionManager.fetchMessages(groupId, Uuid.randomUuid().toString(), FETCH_PARAMS, 0,
                  MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
          // Verify that the fetch request is completed.
          TestUtils.waitForCondition(
              future::isDone,
              DELAYED_SHARE_FETCH_TIMEOUT_MS,
              () -> "Processing in delayed share fetch queue never ended.");
          assertFalse(pendingInitializationFuture1.isDone());
          assertFalse(pendingInitializationFuture2.isDone());
          assertEquals(0, shareGroupMetrics.partitionLoadTimeMs().count());
          // Complete the first pending initialization future.
          pendingInitializationFuture1.complete(null);
          // Verify the partition load time metrics for first partition.
          assertEquals(1, shareGroupMetrics.partitionLoadTimeMs().count());
          assertEquals(90.0, shareGroupMetrics.partitionLoadTimeMs().min());
          assertEquals(90.0, shareGroupMetrics.partitionLoadTimeMs().max());
          assertEquals(90.0, shareGroupMetrics.partitionLoadTimeMs().sum());
          // Complete the second pending initialization future.
          pendingInitializationFuture2.complete(null);
          // Verify the partition load time metrics for both partitions.
          assertEquals(2, shareGroupMetrics.partitionLoadTimeMs().count());
          assertEquals(60.0, shareGroupMetrics.partitionLoadTimeMs().min());
          assertEquals(90.0, shareGroupMetrics.partitionLoadTimeMs().max());
          assertEquals(150.0, shareGroupMetrics.partitionLoadTimeMs().sum());
          shareGroupMetrics.close();
      }
  
      @Test
      public void testDelayedInitializationShouldCompleteFetchRequest() {
          String groupId = "grp";
          Uuid memberId = Uuid.randomUuid();
          Uuid fooId = Uuid.randomUuid();
          TopicIdPartition tp0 = new TopicIdPartition(fooId, new TopicPartition("foo", 0));
          List<TopicIdPartition> topicIdPartitions = List.of(tp0);
  
          SharePartition sp0 = mock(SharePartition.class);
          SharePartitionCache partitionCache = new SharePartitionCache();
          partitionCache.put(new SharePartitionKey(groupId, tp0), sp0);
  
          // Keep the 2 initialization futures pending and 1 completed with leader not available exception.
          CompletableFuture<Void> pendingInitializationFuture1 = new CompletableFuture<>();
          CompletableFuture<Void> pendingInitializationFuture2 = new CompletableFuture<>();
          when(sp0.maybeInitialize()).
              thenReturn(pendingInitializationFuture1)
              .thenReturn(pendingInitializationFuture2)
              .thenReturn(CompletableFuture.failedFuture(new LeaderNotAvailableException("Leader not available")));
  
          Timer mockTimer = systemTimerReaper();
          DelayedOperationPurgatory<DelayedShareFetch> shareFetchPurgatorySpy = spy(new DelayedOperationPurgatory<>(
              "TestShareFetch", mockTimer, mockReplicaManager.localBrokerId(),
              DELAYED_SHARE_FETCH_PURGATORY_PURGE_INTERVAL, false, true));
          mockReplicaManagerDelayedShareFetch(mockReplicaManager, shareFetchPurgatorySpy);
  
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withPartitionCache(partitionCache)
              .withReplicaManager(mockReplicaManager)
              .withTimer(mockTimer)
              .withBrokerTopicStats(brokerTopicStats)
              .build();
  
          // Send 3 requests for share fetch for same share partition.
          CompletableFuture<Map<TopicIdPartition, ShareFetchResponseData.PartitionData>> future1 =
              sharePartitionManager.fetchMessages(groupId, memberId.toString(), FETCH_PARAMS, 0,
                  MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
  
          CompletableFuture<Map<TopicIdPartition, ShareFetchResponseData.PartitionData>> future2 =
              sharePartitionManager.fetchMessages(groupId, memberId.toString(), FETCH_PARAMS, 0,
                  MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
  
          CompletableFuture<Map<TopicIdPartition, ShareFetchResponseData.PartitionData>> future3 =
              sharePartitionManager.fetchMessages(groupId, memberId.toString(), FETCH_PARAMS, 0,
                  MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
  
          Mockito.verify(sp0, times(3)).maybeInitialize();
          Mockito.verify(mockReplicaManager, times(3)).addDelayedShareFetchRequest(any(), any());
          Mockito.verify(shareFetchPurgatorySpy, times(3)).tryCompleteElseWatch(any(), any());
          Mockito.verify(shareFetchPurgatorySpy, times(0)).checkAndComplete(any());
  
          // All 3 requests should be pending.
          assertFalse(future1.isDone());
          assertFalse(future2.isDone());
          assertFalse(future3.isDone());
  
          // Complete one pending initialization future.
          pendingInitializationFuture1.complete(null);
          Mockito.verify(mockReplicaManager, times(1)).completeDelayedShareFetchRequest(any());
          Mockito.verify(shareFetchPurgatorySpy, times(1)).checkAndComplete(any());
  
          pendingInitializationFuture2.complete(null);
          Mockito.verify(mockReplicaManager, times(2)).completeDelayedShareFetchRequest(any());
          Mockito.verify(shareFetchPurgatorySpy, times(2)).checkAndComplete(any());
  
          // Verify that replica manager fetch is not called.
          Mockito.verify(mockReplicaManager, times(0)).readFromLog(
              any(), any(), any(ReplicaQuota.class), anyBoolean());
          // Should have 3 fetch recorded.
          assertEquals(3, brokerTopicStats.allTopicsStats().totalShareFetchRequestRate().count());
          assertEquals(1, brokerTopicStats.numTopics());
          assertEquals(3, brokerTopicStats.topicStats(tp0.topic()).totalShareFetchRequestRate().count());
      }
  
      @Test
      public void testSharePartitionInitializationExceptions() throws Exception {
          String groupId = "grp";
          Uuid memberId = Uuid.randomUuid();
          Uuid fooId = Uuid.randomUuid();
          TopicIdPartition tp0 = new TopicIdPartition(fooId, new TopicPartition("foo", 0));
          List<TopicIdPartition> topicIdPartitions = List.of(tp0);
  
          SharePartition sp0 = mock(SharePartition.class);
          SharePartitionCache partitionCache = new SharePartitionCache();
          partitionCache.put(new SharePartitionKey(groupId, tp0), sp0);
  
          Timer mockTimer = systemTimerReaper();
          DelayedOperationPurgatory<DelayedShareFetch> delayedShareFetchPurgatory = new DelayedOperationPurgatory<>(
                  "TestShareFetch", mockTimer, mockReplicaManager.localBrokerId(),
                  DELAYED_SHARE_FETCH_PURGATORY_PURGE_INTERVAL, false, true);
          mockReplicaManagerDelayedShareFetch(mockReplicaManager, delayedShareFetchPurgatory);
  
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withPartitionCache(partitionCache)
              .withReplicaManager(mockReplicaManager)
              .withTimer(mockTimer)
              .withBrokerTopicStats(brokerTopicStats)
              .build();
  
          // Return LeaderNotAvailableException to simulate initialization failure.
          when(sp0.maybeInitialize()).thenReturn(FutureUtils.failedFuture(new LeaderNotAvailableException("Leader not available")));
          CompletableFuture<Map<TopicIdPartition, ShareFetchResponseData.PartitionData>> future =
              sharePartitionManager.fetchMessages(groupId, memberId.toString(), FETCH_PARAMS, 0,
                  MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
          TestUtils.waitForCondition(
              future::isDone,
              DELAYED_SHARE_FETCH_TIMEOUT_MS,
              () -> "Processing in delayed share fetch queue never ended.");
          // Exception for client should not occur for LeaderNotAvailableException, this exception is to communicate
          // between SharePartitionManager and SharePartition to retry the request as SharePartition is not yet ready.
          assertFalse(future.isCompletedExceptionally());
          assertTrue(future.join().isEmpty());
          Mockito.verify(sp0, times(0)).markFenced();
          // Verify that the share partition is still in the cache on LeaderNotAvailableException.
          assertEquals(1, partitionCache.size());
  
          // Return IllegalStateException to simulate initialization failure.
          when(sp0.maybeInitialize()).thenReturn(FutureUtils.failedFuture(new IllegalStateException("Illegal state")));
          future = sharePartitionManager.fetchMessages(groupId, memberId.toString(), FETCH_PARAMS, 0,
              MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
          TestUtils.waitForCondition(
              future::isDone,
              DELAYED_SHARE_FETCH_TIMEOUT_MS,
              () -> "Processing in delayed share fetch queue never ended.");
          validateShareFetchFutureException(future, tp0, Errors.UNKNOWN_SERVER_ERROR, "Illegal state");
          Mockito.verify(sp0, times(1)).markFenced();
          assertTrue(partitionCache.isEmpty());
  
          // The last exception removes the share partition from the cache hence re-add the share partition to cache.
          partitionCache.put(new SharePartitionKey(groupId, tp0), sp0);
          // Return CoordinatorNotAvailableException to simulate initialization failure.
          when(sp0.maybeInitialize()).thenReturn(FutureUtils.failedFuture(new CoordinatorNotAvailableException("Coordinator not available")));
          future = sharePartitionManager.fetchMessages(groupId, memberId.toString(), FETCH_PARAMS, 0,
              MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
          TestUtils.waitForCondition(
              future::isDone,
              DELAYED_SHARE_FETCH_TIMEOUT_MS,
              () -> "Processing in delayed share fetch queue never ended.");
          validateShareFetchFutureException(future, tp0, Errors.COORDINATOR_NOT_AVAILABLE, "Coordinator not available");
          Mockito.verify(sp0, times(2)).markFenced();
          assertTrue(partitionCache.isEmpty());
  
          // The last exception removes the share partition from the cache hence re-add the share partition to cache.
          partitionCache.put(new SharePartitionKey(groupId, tp0), sp0);
          // Return InvalidRequestException to simulate initialization failure.
          when(sp0.maybeInitialize()).thenReturn(FutureUtils.failedFuture(new InvalidRequestException("Invalid request")));
          future = sharePartitionManager.fetchMessages(groupId, memberId.toString(), FETCH_PARAMS, 0,
              MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
          TestUtils.waitForCondition(
              future::isDone,
              DELAYED_SHARE_FETCH_TIMEOUT_MS,
              () -> "Processing in delayed share fetch queue never ended.");
          validateShareFetchFutureException(future, tp0, Errors.INVALID_REQUEST, "Invalid request");
          Mockito.verify(sp0, times(3)).markFenced();
          assertTrue(partitionCache.isEmpty());
  
          // The last exception removes the share partition from the cache hence re-add the share partition to cache.
          partitionCache.put(new SharePartitionKey(groupId, tp0), sp0);
          // Return FencedStateEpochException to simulate initialization failure.
          when(sp0.maybeInitialize()).thenReturn(FutureUtils.failedFuture(new FencedStateEpochException("Fenced state epoch")));
          future = sharePartitionManager.fetchMessages(groupId, memberId.toString(), FETCH_PARAMS, 0,
              MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
          TestUtils.waitForCondition(
              future::isDone,
              DELAYED_SHARE_FETCH_TIMEOUT_MS,
              () -> "Processing in delayed share fetch queue never ended.");
          validateShareFetchFutureException(future, tp0, Errors.FENCED_STATE_EPOCH, "Fenced state epoch");
          Mockito.verify(sp0, times(4)).markFenced();
          assertTrue(partitionCache.isEmpty());
  
          // The last exception removes the share partition from the cache hence re-add the share partition to cache.
          partitionCache.put(new SharePartitionKey(groupId, tp0), sp0);
          // Return NotLeaderOrFollowerException to simulate initialization failure.
          when(sp0.maybeInitialize()).thenReturn(FutureUtils.failedFuture(new NotLeaderOrFollowerException("Not leader or follower")));
          future = sharePartitionManager.fetchMessages(groupId, memberId.toString(), FETCH_PARAMS, 0,
              MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
          TestUtils.waitForCondition(
              future::isDone,
              DELAYED_SHARE_FETCH_TIMEOUT_MS,
              () -> "Processing in delayed share fetch queue never ended.");
          validateShareFetchFutureException(future, tp0, Errors.NOT_LEADER_OR_FOLLOWER, "Not leader or follower");
          Mockito.verify(sp0, times(5)).markFenced();
          assertTrue(partitionCache.isEmpty());
  
          // The last exception removes the share partition from the cache hence re-add the share partition to cache.
          partitionCache.put(new SharePartitionKey(groupId, tp0), sp0);
          // Return RuntimeException to simulate initialization failure.
          when(sp0.maybeInitialize()).thenReturn(FutureUtils.failedFuture(new RuntimeException("Runtime exception")));
          future = sharePartitionManager.fetchMessages(groupId, memberId.toString(), FETCH_PARAMS, 0,
              MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
          TestUtils.waitForCondition(
              future::isDone,
              DELAYED_SHARE_FETCH_TIMEOUT_MS,
              () -> "Processing in delayed share fetch queue never ended.");
          validateShareFetchFutureException(future, tp0, Errors.UNKNOWN_SERVER_ERROR, "Runtime exception");
          Mockito.verify(sp0, times(6)).markFenced();
          assertTrue(partitionCache.isEmpty());
          // Should have 7 fetch recorded and 6 failures as 1 fetch was waiting on initialization and
          // didn't error out.
          validateBrokerTopicStatsMetrics(
              brokerTopicStats,
              new TopicMetrics(7, 6, 0, 0),
              Map.of(tp0.topic(), new TopicMetrics(7, 6, 0, 0))
          );
      }
  
      @Test
      @SuppressWarnings("unchecked")
      public void testShareFetchProcessingExceptions() throws Exception {
          String groupId = "grp";
          TopicIdPartition tp0 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 0));
          List<TopicIdPartition> topicIdPartitions = List.of(tp0);
  
          SharePartitionCache partitionCache = mock(SharePartitionCache.class);
          // Throw the exception for first fetch request. Return share partition for next.
          when(partitionCache.computeIfAbsent(any(), any()))
              .thenThrow(new RuntimeException("Error creating instance"));
  
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withPartitionCache(partitionCache)
              .withBrokerTopicStats(brokerTopicStats)
              .build();
  
          CompletableFuture<Map<TopicIdPartition, ShareFetchResponseData.PartitionData>> future =
              sharePartitionManager.fetchMessages(groupId, Uuid.randomUuid().toString(), FETCH_PARAMS, 0,
                  MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
          TestUtils.waitForCondition(
              future::isDone,
              DELAYED_SHARE_FETCH_TIMEOUT_MS,
              () -> "Processing for delayed share fetch request not finished.");
          validateShareFetchFutureException(future, tp0, Errors.UNKNOWN_SERVER_ERROR, "Error creating instance");
          // Should have 1 fetch recorded and 1 failure.
          validateBrokerTopicStatsMetrics(
              brokerTopicStats,
              new TopicMetrics(1, 1, 0, 0),
              Map.of(tp0.topic(), new TopicMetrics(1, 1, 0, 0))
          );
      }
  
      @Test
      public void testSharePartitionInitializationFailure() throws Exception {
          String groupId = "grp";
          TopicIdPartition tp0 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 0));
          List<TopicIdPartition> topicIdPartitions = List.of(tp0);
  
          // Send map to check no share partition is created.
          SharePartitionCache partitionCache = new SharePartitionCache();
          // Validate when partition is not the leader.
          Partition partition = mock(Partition.class);
          when(partition.isLeader()).thenReturn(false);
  
          ReplicaManager replicaManager = mock(ReplicaManager.class);
          // First check should throw KafkaStorageException, second check should return partition which
          // is not leader.
          when(replicaManager.getPartitionOrException(any(TopicPartition.class)))
              .thenThrow(new KafkaStorageException("Exception"))
              .thenReturn(partition);
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withReplicaManager(replicaManager)
              .withPartitionCache(partitionCache)
              .withBrokerTopicStats(brokerTopicStats)
              .build();
  
          // Validate when exception is thrown.
          CompletableFuture<Map<TopicIdPartition, ShareFetchResponseData.PartitionData>> future =
              sharePartitionManager.fetchMessages(groupId, Uuid.randomUuid().toString(), FETCH_PARAMS, 0,
                  MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
          TestUtils.waitForCondition(
              future::isDone,
              DELAYED_SHARE_FETCH_TIMEOUT_MS,
              () -> "Processing for delayed share fetch request not finished.");
          validateShareFetchFutureException(future, tp0, Errors.KAFKA_STORAGE_ERROR, "Exception");
          assertTrue(partitionCache.isEmpty());
  
          // Validate when partition is not leader.
          future = sharePartitionManager.fetchMessages(groupId, Uuid.randomUuid().toString(), FETCH_PARAMS, 0,
              MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
          TestUtils.waitForCondition(
              future::isDone,
              DELAYED_SHARE_FETCH_TIMEOUT_MS,
              () -> "Processing for delayed share fetch request not finished.");
          validateShareFetchFutureException(future, tp0, Errors.NOT_LEADER_OR_FOLLOWER);
          assertTrue(partitionCache.isEmpty());
          // Should have 2 fetch recorded and 2 failure.
          validateBrokerTopicStatsMetrics(
              brokerTopicStats,
              new TopicMetrics(2, 2, 0, 0),
              Map.of(tp0.topic(), new TopicMetrics(2, 2, 0, 0))
          );
      }
  
      @Test
      public void testSharePartitionPartialInitializationFailure() throws Exception {
          String groupId = "grp";
          Uuid memberId1 = Uuid.randomUuid();
          // For tp0, share partition instantiation will fail.
          TopicIdPartition tp0 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 0));
          // For tp1, share fetch should succeed.
          TopicIdPartition tp1 = new TopicIdPartition(memberId1, new TopicPartition("foo", 1));
          // For tp2, share partition initialization will fail.
          TopicIdPartition tp2 = new TopicIdPartition(memberId1, new TopicPartition("foo", 2));
          List<TopicIdPartition> topicIdPartitions = List.of(tp0, tp1, tp2);
  
          // Mark partition0 as not the leader.
          Partition partition0 = mock(Partition.class);
          when(partition0.isLeader()).thenReturn(false);
          ReplicaManager replicaManager = mock(ReplicaManager.class);
          when(replicaManager.getPartitionOrException(any(TopicPartition.class)))
              .thenReturn(partition0);
  
          // Mock share partition for tp1, so it can succeed.
          SharePartition sp1 = mock(SharePartition.class);
          SharePartitionCache partitionCache = new SharePartitionCache();
          partitionCache.put(new SharePartitionKey(groupId, tp1), sp1);
  
          when(sp1.maybeAcquireFetchLock(any())).thenReturn(true);
          when(sp1.canAcquireRecords()).thenReturn(true);
          when(sp1.maybeInitialize()).thenReturn(CompletableFuture.completedFuture(null));
          when(sp1.acquire(anyString(), anyInt(), anyInt(), anyLong(), any(), any())).thenReturn(new ShareAcquiredRecords(EMPTY_ACQUIRED_RECORDS, 0));
  
          // Fail initialization for tp2.
          SharePartition sp2 = mock(SharePartition.class);
          partitionCache.put(new SharePartitionKey(groupId, tp2), sp2);
          when(sp2.maybeInitialize()).thenReturn(CompletableFuture.failedFuture(new FencedStateEpochException("Fenced state epoch")));
  
          Timer mockTimer = systemTimerReaper();
          DelayedOperationPurgatory<DelayedShareFetch> delayedShareFetchPurgatory = new DelayedOperationPurgatory<>(
              "TestShareFetch", mockTimer, replicaManager.localBrokerId(),
              DELAYED_SHARE_FETCH_PURGATORY_PURGE_INTERVAL, false, true);
          mockReplicaManagerDelayedShareFetch(replicaManager, delayedShareFetchPurgatory);
          when(sp1.fetchOffsetMetadata(anyLong())).thenReturn(Optional.of(new LogOffsetMetadata(0, 1, 0)));
          mockTopicIdPartitionToReturnDataEqualToMinBytes(replicaManager, tp1, 1);
  
          doAnswer(invocation -> buildLogReadResult(List.of(tp1))).when(replicaManager).readFromLog(any(), any(), any(ReplicaQuota.class), anyBoolean());
  
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withReplicaManager(replicaManager)
              .withPartitionCache(partitionCache)
              .withBrokerTopicStats(brokerTopicStats)
              .withTimer(mockTimer)
              .build();
  
          // Validate when exception is thrown.
          CompletableFuture<Map<TopicIdPartition, ShareFetchResponseData.PartitionData>> future =
              sharePartitionManager.fetchMessages(groupId, Uuid.randomUuid().toString(), FETCH_PARAMS, 0,
                  MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
          assertTrue(future.isDone());
          assertFalse(future.isCompletedExceptionally());
  
          Map<TopicIdPartition, PartitionData> partitionDataMap = future.get();
          assertEquals(3, partitionDataMap.size());
          assertTrue(partitionDataMap.containsKey(tp0));
          assertEquals(Errors.NOT_LEADER_OR_FOLLOWER.code(), partitionDataMap.get(tp0).errorCode());
          assertTrue(partitionDataMap.containsKey(tp1));
          assertEquals(Errors.NONE.code(), partitionDataMap.get(tp1).errorCode());
          assertTrue(partitionDataMap.containsKey(tp2));
          assertEquals(Errors.FENCED_STATE_EPOCH.code(), partitionDataMap.get(tp2).errorCode());
          assertEquals("Fenced state epoch", partitionDataMap.get(tp2).errorMessage());
  
          Mockito.verify(replicaManager, times(0)).completeDelayedShareFetchRequest(any());
          Mockito.verify(replicaManager, times(1)).readFromLog(
              any(), any(), any(ReplicaQuota.class), anyBoolean());
          // Should have 1 fetch recorded and 1 failure as single topic has multiple partition fetch
          // and failure.
          validateBrokerTopicStatsMetrics(
              brokerTopicStats,
              new TopicMetrics(1, 1, 0, 0),
              Map.of(tp0.topic(), new TopicMetrics(1, 1, 0, 0))
          );
      }
  
      @Test
      public void testReplicaManagerFetchException() {
          String groupId = "grp";
          Uuid memberId = Uuid.randomUuid();
          TopicIdPartition tp0 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 0));
          List<TopicIdPartition> topicIdPartitions = List.of(tp0);
  
          SharePartition sp0 = mock(SharePartition.class);
          when(sp0.maybeAcquireFetchLock(any())).thenReturn(true);
          when(sp0.canAcquireRecords()).thenReturn(true);
          when(sp0.maybeInitialize()).thenReturn(CompletableFuture.completedFuture(null));
          SharePartitionCache partitionCache = new SharePartitionCache();
          partitionCache.put(new SharePartitionKey(groupId, tp0), sp0);
  
          Timer mockTimer = systemTimerReaper();
          DelayedOperationPurgatory<DelayedShareFetch> delayedShareFetchPurgatory = new DelayedOperationPurgatory<>(
              "TestShareFetch", mockTimer, mockReplicaManager.localBrokerId(),
              DELAYED_SHARE_FETCH_PURGATORY_PURGE_INTERVAL, false, true);
          mockReplicaManagerDelayedShareFetch(mockReplicaManager, delayedShareFetchPurgatory);
  
          doThrow(new RuntimeException("Exception")).when(mockReplicaManager).readFromLog(any(), any(), any(ReplicaQuota.class), anyBoolean());
  
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withPartitionCache(partitionCache)
              .withReplicaManager(mockReplicaManager)
              .withTimer(mockTimer)
              .withBrokerTopicStats(brokerTopicStats)
              .build();
  
          CompletableFuture<Map<TopicIdPartition, ShareFetchResponseData.PartitionData>> future =
              sharePartitionManager.fetchMessages(groupId, memberId.toString(), FETCH_PARAMS, 0,
                  MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
          validateShareFetchFutureException(future, tp0, Errors.UNKNOWN_SERVER_ERROR, "Exception");
          // Verify that the share partition is still in the cache on exception.
          assertEquals(1, partitionCache.size());
  
          // Throw NotLeaderOrFollowerException from replica manager fetch which should evict instance from the cache.
          doThrow(new NotLeaderOrFollowerException("Leader exception")).when(mockReplicaManager).readFromLog(any(), any(), any(ReplicaQuota.class), anyBoolean());
  
          future = sharePartitionManager.fetchMessages(groupId, memberId.toString(), FETCH_PARAMS, 0,
              MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
          validateShareFetchFutureException(future, tp0, Errors.NOT_LEADER_OR_FOLLOWER, "Leader exception");
          assertTrue(partitionCache.isEmpty());
          // Should have 2 fetch recorded and 2 failures.
          validateBrokerTopicStatsMetrics(
              brokerTopicStats,
              new TopicMetrics(2, 2, 0, 0),
              Map.of(tp0.topic(), new TopicMetrics(2, 2, 0, 0))
          );
      }
  
      @Test
      public void testReplicaManagerFetchMultipleSharePartitionsException() {
          String groupId = "grp";
          Uuid memberId = Uuid.randomUuid();
  
          TopicIdPartition tp0 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 0));
          TopicIdPartition tp1 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("bar", 0));
          List<TopicIdPartition> topicIdPartitions = List.of(tp0, tp1);
  
          SharePartition sp0 = mock(SharePartition.class);
          when(sp0.maybeAcquireFetchLock(any())).thenReturn(true);
          when(sp0.canAcquireRecords()).thenReturn(true);
          when(sp0.maybeInitialize()).thenReturn(CompletableFuture.completedFuture(null));
  
          SharePartition sp1 = mock(SharePartition.class);
          // Do not make the share partition acquirable hence it shouldn't be removed from the cache,
          // as it won't be part of replica manager readFromLog request.
          when(sp1.maybeAcquireFetchLock(any())).thenReturn(false);
          when(sp1.maybeInitialize()).thenReturn(CompletableFuture.completedFuture(null));
  
          SharePartitionCache partitionCache = new SharePartitionCache();
          partitionCache.put(new SharePartitionKey(groupId, tp0), sp0);
          partitionCache.put(new SharePartitionKey(groupId, tp1), sp1);
  
          Timer mockTimer = systemTimerReaper();
          DelayedOperationPurgatory<DelayedShareFetch> delayedShareFetchPurgatory = new DelayedOperationPurgatory<>(
              "TestShareFetch", mockTimer, mockReplicaManager.localBrokerId(),
              DELAYED_SHARE_FETCH_PURGATORY_PURGE_INTERVAL, false, true);
          mockReplicaManagerDelayedShareFetch(mockReplicaManager, delayedShareFetchPurgatory);
  
          // Throw FencedStateEpochException from replica manager fetch which should evict instance from the cache.
          doThrow(new FencedStateEpochException("Fenced exception")).when(mockReplicaManager).readFromLog(any(), any(), any(ReplicaQuota.class), anyBoolean());
  
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withPartitionCache(partitionCache)
              .withReplicaManager(mockReplicaManager)
              .withTimer(mockTimer)
              .withBrokerTopicStats(brokerTopicStats)
              .build();
  
          CompletableFuture<Map<TopicIdPartition, ShareFetchResponseData.PartitionData>> future =
              sharePartitionManager.fetchMessages(groupId, memberId.toString(), FETCH_PARAMS, 0,
                  MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
          validateShareFetchFutureException(future, tp0, Errors.FENCED_STATE_EPOCH, "Fenced exception");
          // Verify that tp1 is still in the cache on exception.
          assertEquals(1, partitionCache.size());
          assertEquals(sp1, partitionCache.get(new SharePartitionKey(groupId, tp1)));
  
          // Make sp1 acquirable and add sp0 back in partition cache. Both share partitions should be
          // removed from the cache.
          when(sp1.maybeAcquireFetchLock(any())).thenReturn(true);
          when(sp1.canAcquireRecords()).thenReturn(true);
          partitionCache.put(new SharePartitionKey(groupId, tp0), sp0);
          // Throw FencedStateEpochException from replica manager fetch which should evict instance from the cache.
          doThrow(new FencedStateEpochException("Fenced exception again")).when(mockReplicaManager).readFromLog(any(), any(), any(ReplicaQuota.class), anyBoolean());
  
          future = sharePartitionManager.fetchMessages(groupId, memberId.toString(), FETCH_PARAMS, 0,
              MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
          validateShareFetchFutureException(future, List.of(tp0, tp1), Errors.FENCED_STATE_EPOCH, "Fenced exception again");
          assertTrue(partitionCache.isEmpty());
          // Should have 4 fetch recorded (2 fetch and 2 topics) and 3 failures as sp1 was not acquired
          // in first fetch and shall have empty response. Similarly, tp0 should record 2 failures and
          // tp1 should record 1 failure.
          validateBrokerTopicStatsMetrics(
              brokerTopicStats,
              new TopicMetrics(4, 3, 0, 0),
              Map.of(tp0.topic(), new TopicMetrics(2, 2, 0, 0), tp1.topic(), new TopicMetrics(2, 1, 0, 0))
          );
      }
  
      @Test
      public void testListenerRegistration() {
          String groupId = "grp";
          Uuid memberId = Uuid.randomUuid();
  
          TopicIdPartition tp0 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 0));
          TopicIdPartition tp1 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("bar", 0));
          List<TopicIdPartition> topicIdPartitions = List.of(tp0, tp1);
  
          ReplicaManager mockReplicaManager = mock(ReplicaManager.class);
          Partition partition = mockPartition();
          when(mockReplicaManager.getPartitionOrException((TopicPartition) Mockito.any())).thenReturn(partition);
  
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withReplicaManager(mockReplicaManager)
              .withBrokerTopicStats(brokerTopicStats)
              .build();
  
          CompletableFuture<Map<TopicIdPartition, PartitionData>> future = sharePartitionManager.fetchMessages(
              groupId, memberId.toString(), FETCH_PARAMS, 0, MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
          assertTrue(future.isDone());
          // Validate that the listener is registered.
          verify(mockReplicaManager, times(2)).maybeAddListener(any(), any());
          // The share partition initialization should error out as further mocks are not provided, the
          // metrics should mark fetch as failed.
          validateBrokerTopicStatsMetrics(
              brokerTopicStats,
              new TopicMetrics(2, 2, 0, 0),
              Map.of(tp0.topic(), new TopicMetrics(1, 1, 0, 0), tp1.topic(), new TopicMetrics(1, 1, 0, 0))
          );
      }
  
      @Test
      public void testSharePartitionListenerOnFailed() {
          SharePartitionKey sharePartitionKey = new SharePartitionKey("grp",
              new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 0)));
          SharePartitionCache partitionCache = new SharePartitionCache();
          ReplicaManager mockReplicaManager = mock(ReplicaManager.class);
  
          SharePartitionListener partitionListener = new SharePartitionListener(sharePartitionKey, mockReplicaManager, partitionCache);
          testSharePartitionListener(sharePartitionKey, partitionCache, mockReplicaManager, partitionListener::onFailed);
      }
  
      @Test
      public void testSharePartitionListenerOnDeleted() {
          SharePartitionKey sharePartitionKey = new SharePartitionKey("grp",
              new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 0)));
          SharePartitionCache partitionCache = new SharePartitionCache();
          ReplicaManager mockReplicaManager = mock(ReplicaManager.class);
  
          SharePartitionListener partitionListener = new SharePartitionListener(sharePartitionKey, mockReplicaManager, partitionCache);
          testSharePartitionListener(sharePartitionKey, partitionCache, mockReplicaManager, partitionListener::onDeleted);
      }
  
      @Test
      public void testSharePartitionListenerOnBecomingFollower() {
          SharePartitionKey sharePartitionKey = new SharePartitionKey("grp",
              new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 0)));
          SharePartitionCache partitionCache = new SharePartitionCache();
          ReplicaManager mockReplicaManager = mock(ReplicaManager.class);
  
          SharePartitionListener partitionListener = new SharePartitionListener(sharePartitionKey, mockReplicaManager, partitionCache);
          testSharePartitionListener(sharePartitionKey, partitionCache, mockReplicaManager, partitionListener::onBecomingFollower);
      }
  
      @Test
      public void testFetchMessagesRotatePartitions() {
          String groupId = "grp";
          Uuid memberId1 = Uuid.randomUuid();
          TopicIdPartition tp0 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 0));
          TopicIdPartition tp1 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 1));
          TopicIdPartition tp2 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("bar", 0));
          TopicIdPartition tp3 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("bar", 1));
          TopicIdPartition tp4 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 2));
          TopicIdPartition tp5 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("bar", 2));
          TopicIdPartition tp6 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 3));
          List<TopicIdPartition> topicIdPartitions = List.of(tp0, tp1, tp2, tp3, tp4, tp5, tp6);
  
          sharePartitionManager = Mockito.spy(SharePartitionManagerBuilder.builder().withBrokerTopicStats(brokerTopicStats).build());
          // Capture the arguments passed to processShareFetch.
          ArgumentCaptor<ShareFetch> captor = ArgumentCaptor.forClass(ShareFetch.class);
  
          sharePartitionManager.fetchMessages(groupId, memberId1.toString(), FETCH_PARAMS, 0,
              MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
          verify(sharePartitionManager, times(1)).processShareFetch(captor.capture());
          // Verify the partitions rotation, no rotation.
          ShareFetch resultShareFetch = captor.getValue();
          validateRotatedListEquals(resultShareFetch.topicIdPartitions(), topicIdPartitions, 0);
  
          // Single rotation.
          sharePartitionManager.fetchMessages(groupId, memberId1.toString(), FETCH_PARAMS, 1,
              MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
          verify(sharePartitionManager, times(2)).processShareFetch(captor.capture());
          // Verify the partitions rotation, rotate by 1.
          resultShareFetch = captor.getValue();
          validateRotatedListEquals(topicIdPartitions, resultShareFetch.topicIdPartitions(), 1);
  
          // Rotation by 3, less that the number of partitions.
          sharePartitionManager.fetchMessages(groupId, memberId1.toString(), FETCH_PARAMS, 3,
              MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
          verify(sharePartitionManager, times(3)).processShareFetch(captor.capture());
          // Verify the partitions rotation, rotate by 3.
          resultShareFetch = captor.getValue();
          validateRotatedListEquals(topicIdPartitions, resultShareFetch.topicIdPartitions(), 3);
  
          // Rotation by 12, more than the number of partitions.
          sharePartitionManager.fetchMessages(groupId, memberId1.toString(), FETCH_PARAMS, 12,
              MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
          verify(sharePartitionManager, times(4)).processShareFetch(captor.capture());
          // Verify the partitions rotation, rotate by 5 (12 % 7).
          resultShareFetch = captor.getValue();
          validateRotatedListEquals(topicIdPartitions, resultShareFetch.topicIdPartitions(), 5);
          // Rotation by Integer.MAX_VALUE, boundary test.
          sharePartitionManager.fetchMessages(groupId, memberId1.toString(), FETCH_PARAMS, Integer.MAX_VALUE,
              MAX_FETCH_RECORDS, BATCH_SIZE, topicIdPartitions);
          verify(sharePartitionManager, times(5)).processShareFetch(captor.capture());
          // Verify the partitions rotation, rotate by 1 (2147483647 % 7).
          resultShareFetch = captor.getValue();
          validateRotatedListEquals(topicIdPartitions, resultShareFetch.topicIdPartitions(), 1);
      }
  
      @Test
      public void testCreateIdleShareFetchTask() throws Exception {
          ReplicaManager replicaManager = mock(ReplicaManager.class);
  
          MockTimer mockTimer = new MockTimer(time);
          long maxWaitMs = 1000L;
  
          // Set up the mock to capture and add the timer task
          Mockito.doAnswer(invocation -> {
              TimerTask timerTask = invocation.getArgument(0);
              mockTimer.add(timerTask);
              return null;
          }).when(replicaManager).addShareFetchTimerRequest(Mockito.any(TimerTask.class));
  
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withReplicaManager(replicaManager)
              .withTime(time)
              .withTimer(mockTimer)
              .build();
  
          CompletableFuture<Void> future = sharePartitionManager.createIdleShareFetchTimerTask(maxWaitMs);
          // Future should not be completed immediately
          assertFalse(future.isDone());
  
          mockTimer.advanceClock(maxWaitMs / 2);
          assertFalse(future.isDone());
  
          mockTimer.advanceClock((maxWaitMs / 2) + 1);
          // Verify the future is completed after the wait time
          assertTrue(future.isDone());
          assertFalse(future.isCompletedExceptionally());
      }
  
      @Test
      public void testOnShareVersionToggle() {
          String groupId = "grp";
          SharePartition sp0 = mock(SharePartition.class);
          SharePartition sp1 = mock(SharePartition.class);
          SharePartition sp2 = mock(SharePartition.class);
          SharePartition sp3 = mock(SharePartition.class);
  
          // Mock the share partitions corresponding to the topic partitions.
          SharePartitionCache partitionCache = new SharePartitionCache();
          partitionCache.put(
              new SharePartitionKey(groupId, new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo1", 0))), sp0
          );
          partitionCache.put(
              new SharePartitionKey(groupId, new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo2", 0))), sp1
          );
          partitionCache.put(
              new SharePartitionKey(groupId, new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo3", 0))), sp2
          );
          partitionCache.put(
              new SharePartitionKey(groupId, new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo4", 0))), sp3
          );
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withPartitionCache(partitionCache)
              .build();
          assertEquals(4, partitionCache.size());
          sharePartitionManager.onShareVersionToggle(ShareVersion.SV_0, false);
          // Because we are toggling to a share version which does not support share groups, the cache inside share partitions must be cleared.
          assertEquals(0, partitionCache.size());
          //Check if all share partitions have been fenced.
          Mockito.verify(sp0).markFenced();
          Mockito.verify(sp1).markFenced();
          Mockito.verify(sp2).markFenced();
          Mockito.verify(sp3).markFenced();
      }
  
      @Test
      public void testOnShareVersionToggleWhenEnabledFromConfig() {
          SharePartition sp0 = mock(SharePartition.class);
          // Mock the share partitions corresponding to the topic partitions.
          SharePartitionCache partitionCache = new SharePartitionCache();
          partitionCache.put(
              new SharePartitionKey("grp", new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 0))), sp0
          );
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withPartitionCache(partitionCache)
              .build();
          assertEquals(1, partitionCache.size());
          sharePartitionManager.onShareVersionToggle(ShareVersion.SV_0, true);
          // Though share version is toggled to off, but it's enabled from config, hence the cache should not be cleared.
          assertEquals(1, partitionCache.size());
          Mockito.verify(sp0, times(0)).markFenced();
      }
  
      @Test
      public void testShareGroupListener() {
          String groupId = "grp";
          TopicIdPartition tp0 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 0));
          TopicIdPartition tp1 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 1));
          Uuid memberId1 = Uuid.randomUuid();
          Uuid memberId2 = Uuid.randomUuid();
  
          SharePartition sp0 = mock(SharePartition.class);
          SharePartition sp1 = mock(SharePartition.class);
  
          ShareSessionCache cache = new ShareSessionCache(10);
          cache.maybeCreateSession(groupId, memberId1, new ImplicitLinkedHashCollection<>(), CONNECTION_ID);
          cache.maybeCreateSession(groupId, memberId2, new ImplicitLinkedHashCollection<>(), "id-2");
  
          SharePartitionCache partitionCache = new SharePartitionCache();
          partitionCache.computeIfAbsent(new SharePartitionKey(groupId, tp0), k -> sp0);
          partitionCache.computeIfAbsent(new SharePartitionKey(groupId, tp1), k -> sp1);
  
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withCache(cache)
              .withPartitionCache(partitionCache)
              .withReplicaManager(mockReplicaManager)
              .build();
  
          assertEquals(2, cache.size());
          assertEquals(2, partitionCache.size());
  
          // Invoke listeners by simulating connection disconnect for memberId1.
          cache.connectionDisconnectListener().onDisconnect(CONNECTION_ID);
          // Session cache should remove the memberId1.
          assertEquals(1, cache.size());
          // Partition cache should not remove the share partitions as the group is not empty.
          assertEquals(2, partitionCache.size());
          assertNotNull(cache.get(new ShareSessionKey(groupId, memberId2)));
  
          // Invoke listeners by simulating connection disconnect for memberId2.
          cache.connectionDisconnectListener().onDisconnect("id-2");
          // Session cache should remove the memberId2.
          assertEquals(0, cache.size());
          // Partition cache should remove the share partitions as the group is empty.
          assertEquals(0, partitionCache.size());
  
          Mockito.verify(sp0, times(1)).markFenced();
          Mockito.verify(sp1, times(1)).markFenced();
          Mockito.verify(mockReplicaManager, times(2)).removeListener(any(), any());
      }
  
      @Test
      public void testShareGroupListenerWithEmptyCache() {
          String groupId = "grp";
          TopicIdPartition tp0 = new TopicIdPartition(Uuid.randomUuid(), new TopicPartition("foo", 0));
          Uuid memberId1 = Uuid.randomUuid();
  
          SharePartition sp0 = mock(SharePartition.class);
  
          ShareSessionCache cache = new ShareSessionCache(10);
          cache.maybeCreateSession(groupId, memberId1, new ImplicitLinkedHashCollection<>(), CONNECTION_ID);
  
          SharePartitionCache partitionCache = spy(new SharePartitionCache());
          partitionCache.computeIfAbsent(new SharePartitionKey(groupId, tp0), k -> sp0);
  
          sharePartitionManager = SharePartitionManagerBuilder.builder()
              .withCache(cache)
              .withPartitionCache(partitionCache)
              .withReplicaManager(mockReplicaManager)
              .build();
  
          assertEquals(1, cache.size());
          assertEquals(1, partitionCache.size());
  
          // Clean up share session and partition cache.
          sharePartitionManager.onShareVersionToggle(ShareVersion.SV_0, false);
          assertEquals(0, cache.size());
          assertEquals(0, partitionCache.size());
  
          Mockito.verify(sp0, times(1)).markFenced();
          Mockito.verify(mockReplicaManager, times(1)).removeListener(any(), any());
          Mockito.verify(partitionCache, times(0)).topicIdPartitionsForGroup(groupId);
  
          // Invoke listeners by simulating connection disconnect for member. As the group is empty,
          // hence onGroupEmpty method should be invoked and should complete without any exception.
          cache.connectionDisconnectListener().onDisconnect(CONNECTION_ID);
          // Verify that the listener is called for the group.
          Mockito.verify(partitionCache, times(1)).topicIdPartitionsForGroup(groupId);
      }
  
      private Timer systemTimerReaper() {
          return new SystemTimerReaper(
              TIMER_NAME_PREFIX + "-test-reaper",
              new SystemTimer(TIMER_NAME_PREFIX + "-test-timer"));
      }
  
      private void assertNoReaperThreadsPendingClose() throws InterruptedException {
          TestUtils.waitForCondition(
              () -> Thread.getAllStackTraces().keySet().stream().noneMatch(t -> t.getName().contains(TIMER_NAME_PREFIX)),
              "Found unexpected reaper threads with name containing: " + TIMER_NAME_PREFIX);
      }
  
      private void testSharePartitionListener(
          SharePartitionKey sharePartitionKey,
          SharePartitionCache partitionCache,
          ReplicaManager mockReplicaManager,
          Consumer<TopicPartition> listenerConsumer
      ) {
          // Add another share partition to the cache.
          TopicPartition tp = new TopicPartition("foo", 1);
          TopicIdPartition tpId = new TopicIdPartition(Uuid.randomUuid(), tp);
          SharePartitionKey spk = new SharePartitionKey("grp", tpId);
  
          SharePartition sp0 = mock(SharePartition.class);
          SharePartition sp1 = mock(SharePartition.class);
          partitionCache.put(sharePartitionKey, sp0);
          partitionCache.put(spk, sp1);
  
          // Invoke listener for first share partition.
          listenerConsumer.accept(sharePartitionKey.topicIdPartition().topicPartition());
  
          // Validate that the share partition is removed from the cache.
          assertEquals(1, partitionCache.size());
          assertFalse(partitionCache.containsKey(sharePartitionKey));
          verify(sp0, times(1)).markFenced();
          verify(mockReplicaManager, times(1)).removeListener(any(), any());
  
          // Invoke listener for non-matching share partition.
          listenerConsumer.accept(tp);
          // The non-matching share partition should not be removed as the listener is attached to a different topic partition.
          assertEquals(1, partitionCache.size());
          verify(sp1, times(0)).markFenced();
          // Verify the remove listener is not called for the second share partition.
          verify(mockReplicaManager, times(1)).removeListener(any(), any());
      }
  
      private ShareFetchResponseData.PartitionData noErrorShareFetchResponse() {
          return new ShareFetchResponseData.PartitionData().setPartitionIndex(0);
      }
  
      private ShareFetchResponseData.PartitionData errorShareFetchResponse(Short errorCode) {
          return new ShareFetchResponseData.PartitionData().setPartitionIndex(0).setErrorCode(errorCode);
      }
  
      private void mockUpdateAndGenerateResponseData(ShareFetchContext context, String groupId, Uuid memberId) {
          LinkedHashMap<TopicIdPartition, ShareFetchResponseData.PartitionData> data = new LinkedHashMap<>();
          if (context.getClass() == ShareSessionContext.class) {
              ShareSessionContext shareSessionContext = (ShareSessionContext) context;
              if (!shareSessionContext.isSubsequent()) {
                  shareSessionContext.shareFetchData().forEach(topicIdPartition -> data.put(topicIdPartition,
                          topicIdPartition.topic() == null ? errorShareFetchResponse(Errors.UNKNOWN_TOPIC_ID.code()) :
                                  noErrorShareFetchResponse()));
              } else {
                  synchronized (shareSessionContext.session()) {
                      shareSessionContext.session().partitionMap().forEach(cachedSharePartition -> {
                          TopicIdPartition topicIdPartition = new TopicIdPartition(cachedSharePartition.topicId(),
                              new TopicPartition(cachedSharePartition.topic(), cachedSharePartition.partition()));
                          data.put(topicIdPartition, topicIdPartition.topic() == null ? errorShareFetchResponse(Errors.UNKNOWN_TOPIC_ID.code()) :
                                  noErrorShareFetchResponse());
                      });
                  }
              }
          }
          context.updateAndGenerateResponseData(groupId, memberId, data);
      }
  
      private void assertPartitionsPresent(ShareSessionContext context, List<TopicIdPartition> partitions) {
          Set<TopicIdPartition> partitionsInContext = new HashSet<>();
          if (!context.isSubsequent()) {
              partitionsInContext.addAll(context.shareFetchData());
          } else {
              context.session().partitionMap().forEach(cachedSharePartition -> {
                  TopicIdPartition topicIdPartition = new TopicIdPartition(cachedSharePartition.topicId(), new
                          TopicPartition(cachedSharePartition.topic(), cachedSharePartition.partition()));
                  partitionsInContext.add(topicIdPartition);
              });
          }
          Set<TopicIdPartition> partitionsSet = new HashSet<>(partitions);
          assertEquals(partitionsSet, partitionsInContext);
      }
  
      private void assertErroneousAndValidTopicIdPartitions(
          ErroneousAndValidPartitionData erroneousAndValidPartitionData,
          List<TopicIdPartition> expectedErroneous, List<TopicIdPartition> expectedValid) {
          Set<TopicIdPartition> expectedErroneousSet = new HashSet<>(expectedErroneous);
          Set<TopicIdPartition> expectedValidSet = new HashSet<>(expectedValid);
          Set<TopicIdPartition> actualErroneousPartitions = new HashSet<>();
          erroneousAndValidPartitionData.erroneous().forEach((topicIdPartition, partitionData) ->
                  actualErroneousPartitions.add(topicIdPartition));
          Set<TopicIdPartition> actualValidPartitions = new HashSet<>(erroneousAndValidPartitionData.validTopicIdPartitions());
          assertEquals(expectedErroneousSet, actualErroneousPartitions);
          assertEquals(expectedValidSet, actualValidPartitions);
      }
  
      private Partition mockPartition() {
          Partition partition = mock(Partition.class);
          when(partition.isLeader()).thenReturn(true);
          when(partition.getLeaderEpoch()).thenReturn(1);
  
          return partition;
      }
  
      private void validateShareFetchFutureException(CompletableFuture<Map<TopicIdPartition, PartitionData>> future,
          TopicIdPartition topicIdPartition, Errors error) {
          validateShareFetchFutureException(future, List.of(topicIdPartition), error, null);
      }
  
      private void validateShareFetchFutureException(CompletableFuture<Map<TopicIdPartition, PartitionData>> future,
          TopicIdPartition topicIdPartition, Errors error, String message) {
          validateShareFetchFutureException(future, List.of(topicIdPartition), error, message);
      }
  
      private void validateShareFetchFutureException(CompletableFuture<Map<TopicIdPartition, PartitionData>> future,
          List<TopicIdPartition> topicIdPartitions, Errors error, String message) {
          assertFalse(future.isCompletedExceptionally());
          Map<TopicIdPartition, ShareFetchResponseData.PartitionData> result = future.join();
          assertEquals(topicIdPartitions.size(), result.size());
          topicIdPartitions.forEach(topicIdPartition -> {
              assertTrue(result.containsKey(topicIdPartition));
              assertEquals(topicIdPartition.partition(), result.get(topicIdPartition).partitionIndex());
              assertEquals(error.code(), result.get(topicIdPartition).errorCode());
              assertEquals(message, result.get(topicIdPartition).errorMessage());
          });
      }
  
      private void mockFetchOffsetForTimestamp(ReplicaManager replicaManager) {
          FileRecords.TimestampAndOffset timestampAndOffset = new FileRecords.TimestampAndOffset(-1L, 0L, Optional.empty());
          Mockito.doReturn(new OffsetResultHolder(Optional.of(timestampAndOffset), Optional.empty())).
              when(replicaManager).fetchOffsetForTimestamp(Mockito.any(TopicPartition.class), Mockito.anyLong(), Mockito.any(), Mockito.any(), Mockito.anyBoolean());
      }
  
      private void validateBrokerTopicStatsMetrics(
          BrokerTopicStats brokerTopicStats,
          TopicMetrics expectedAllTopicMetrics,
          Map<String, TopicMetrics> expectedTopicMetrics
      ) {
          if (expectedAllTopicMetrics != null) {
              assertEquals(expectedAllTopicMetrics.totalShareFetchRequestCount, brokerTopicStats.allTopicsStats().totalShareFetchRequestRate().count());
              assertEquals(expectedAllTopicMetrics.failedShareFetchRequestCount, brokerTopicStats.allTopicsStats().failedShareFetchRequestRate().count());
              assertEquals(expectedAllTopicMetrics.totalShareAcknowledgementRequestCount, brokerTopicStats.allTopicsStats().totalShareAcknowledgementRequestRate().count());
              assertEquals(expectedAllTopicMetrics.failedShareAcknowledgementRequestCount, brokerTopicStats.allTopicsStats().failedShareAcknowledgementRequestRate().count());
          }
          // Validate tracked topic metrics.
          assertEquals(expectedTopicMetrics.size(), brokerTopicStats.numTopics());
          expectedTopicMetrics.forEach((topic, metrics) -> {
              BrokerTopicMetrics topicMetrics = brokerTopicStats.topicStats(topic);
              assertEquals(metrics.totalShareFetchRequestCount, topicMetrics.totalShareFetchRequestRate().count());
              assertEquals(metrics.failedShareFetchRequestCount, topicMetrics.failedShareFetchRequestRate().count());
              assertEquals(metrics.totalShareAcknowledgementRequestCount, topicMetrics.totalShareAcknowledgementRequestRate().count());
              assertEquals(metrics.failedShareAcknowledgementRequestCount, topicMetrics.failedShareAcknowledgementRequestRate().count());
          });
      }
  
      static Seq<Tuple2<TopicIdPartition, LogReadResult>> buildLogReadResult(List<TopicIdPartition> topicIdPartitions) {
          List<Tuple2<TopicIdPartition, LogReadResult>> logReadResults = new ArrayList<>();
          topicIdPartitions.forEach(topicIdPartition -> logReadResults.add(new Tuple2<>(topicIdPartition, new LogReadResult(
              new FetchDataInfo(new LogOffsetMetadata(0, 0, 0), MemoryRecords.withRecords(
                      Compression.NONE, new SimpleRecord("test-key".getBytes(), "test-value".getBytes()))),
              Optional.empty(),
              -1L,
              -1L,
              -1L,
              -1L,
              -1L,
              OptionalLong.empty(),
              OptionalInt.empty(),
              Optional.empty()
          ))));
          return CollectionConverters.asScala(logReadResults).toSeq();
      }
  
      static void mockReplicaManagerDelayedShareFetch(ReplicaManager replicaManager,
                                                      DelayedOperationPurgatory<DelayedShareFetch> delayedShareFetchPurgatory) {
          doAnswer(invocationOnMock -> {
              Object[] args = invocationOnMock.getArguments();
              DelayedShareFetchKey key = (DelayedShareFetchKey) args[0];
              delayedShareFetchPurgatory.checkAndComplete(key);
              return null;
          }).when(replicaManager).completeDelayedShareFetchRequest(any(DelayedShareFetchKey.class));
  
          doAnswer(invocationOnMock -> {
              Object[] args = invocationOnMock.getArguments();
              DelayedShareFetch operation = (DelayedShareFetch) args[0];
              List<DelayedOperationKey> keys = (List<DelayedOperationKey>) args[1];
              delayedShareFetchPurgatory.tryCompleteElseWatch(operation, keys);
              return null;
          }).when(replicaManager).addDelayedShareFetchRequest(any(), any());
      }
  
      private record TopicMetrics(
          long totalShareFetchRequestCount,
          long failedShareFetchRequestCount,
          long totalShareAcknowledgementRequestCount,
          long failedShareAcknowledgementRequestCount
      ) { }
  
      static class SharePartitionManagerBuilder {
          private final Persister persister = new NoOpStatePersister();
          private ReplicaManager replicaManager = mock(ReplicaManager.class);
          private Time time = new MockTime();
          private ShareSessionCache cache = new ShareSessionCache(10);
          private SharePartitionCache partitionCache = new SharePartitionCache();
          private Timer timer = new MockTimer();
          private ShareGroupMetrics shareGroupMetrics = new ShareGroupMetrics(time);
          private BrokerTopicStats brokerTopicStats;
  
          private SharePartitionManagerBuilder withReplicaManager(ReplicaManager replicaManager) {
              this.replicaManager = replicaManager;
              return this;
          }
  
          private SharePartitionManagerBuilder withTime(Time time) {
              this.time = time;
              return this;
          }
  
          private SharePartitionManagerBuilder withCache(ShareSessionCache cache) {
              this.cache = cache;
              return this;
          }
  
          SharePartitionManagerBuilder withPartitionCache(SharePartitionCache partitionCache) {
              this.partitionCache = partitionCache;
              return this;
          }
  
          private SharePartitionManagerBuilder withTimer(Timer timer) {
              this.timer = timer;
              return this;
          }
  
          private SharePartitionManagerBuilder withShareGroupMetrics(ShareGroupMetrics shareGroupMetrics) {
              this.shareGroupMetrics = shareGroupMetrics;
              return this;
          }
  
          private SharePartitionManagerBuilder withBrokerTopicStats(BrokerTopicStats brokerTopicStats) {
              this.brokerTopicStats = brokerTopicStats;
              return this;
          }
  
          public static SharePartitionManagerBuilder builder() {
              return new SharePartitionManagerBuilder();
          }
  
          public SharePartitionManager build() {
              return new SharePartitionManager(replicaManager,
                  time,
                  cache,
                  partitionCache,
                  DEFAULT_RECORD_LOCK_DURATION_MS,
                  timer,
                  MAX_DELIVERY_COUNT,
                  MAX_IN_FLIGHT_MESSAGES,
                  REMOTE_FETCH_MAX_WAIT_MS,
                  persister,
                  mock(GroupConfigManager.class),
                  shareGroupMetrics,
                  brokerTopicStats
              );
          }
      }
  }
  