Repository: apache/kafka
File: core/src/test/java/kafka/server/share/DelayedShareFetchTest.java (Java)

Code before:
            future, List.of(tp0, tp1), BATCH_SIZE, MAX_FETCH_RECORDS,
            BROKER_TOPIC_STATS);

        when(sp0.nextFetchOffset()).thenReturn(10L);
        // Fetch offset does not match with the cached entry for sp0. Hence, a replica manager fetch will happen for sp0.
        when(sp0.fetchOffsetMetadata(anyLong())).thenReturn(Optional.empty());

        // Mocking remote storage read result for tp0.
        doAnswer(invocation -> buildLocalAndRemoteFetchResult(Set.of(), Set.of(tp0))).when(replicaManager).readFromLog(any(), any(), any(ReplicaQuota.class), anyBoolean());

        // Remote fetch related mocks. Remote fetch object completes within tryComplete in this mock, hence request will move on to forceComplete.
        RemoteLogReadResult remoteFetchResult = new RemoteLogReadResult(
            Optional.empty(),
            Optional.of(new TimeoutException("Error occurred while creating remote fetch result")) // Remote fetch result is returned with an error.
        );
        RemoteLogManager remoteLogManager = mock(RemoteLogManager.class);
        doAnswer(invocationOnMock -> {
            // Make sure that the callback is called to populate remoteFetchResult for the mock behaviour.
            Consumer<RemoteLogReadResult> callback = invocationOnMock.getArgument(1);
            callback.accept(remoteFetchResult);
            return CompletableFuture.completedFuture(remoteFetchResult);
        }).when(remoteLogManager).asyncRead(any(), any());
        when(replicaManager.remoteLogManager()).thenReturn(Option.apply(remoteLogManager));

        Uuid fetchId = Uuid.randomUuid();

        Partition p0 = mock(Partition.class);
        when(p0.isLeader()).thenReturn(true);

        Partition p1 = mock(Partition.class);
        when(p1.isLeader()).thenReturn(true);

        when(replicaManager.getPartitionOrException(tp0.topicPartition())).thenReturn(p0);
        when(replicaManager.getPartitionOrException(tp1.topicPartition())).thenReturn(p1);

        BiConsumer<SharePartitionKey, Throwable> exceptionHandler = mockExceptionHandler();
        DelayedShareFetch delayedShareFetch = spy(DelayedShareFetchBuilder.builder()
            .withShareFetchData(shareFetch)
            .withSharePartitions(sharePartitions)
            .withReplicaManager(replicaManager)
            .withPartitionMaxBytesStrategy(mockPartitionMaxBytes(Set.of(tp0, tp1)))
            .withFetchId(fetchId)
            .withExceptionHandler(exceptionHandler)
            .build());

        // sp0 is acquirable, sp1 is not acquirable.
        when(sp0.maybeAcquireFetchLock(fetchId)).thenReturn(true);
        when(sp1.maybeAcquireFetchLock(fetchId)).thenReturn(false);

        assertFalse(delayedShareFetch.isCompleted());
        assertTrue(delayedShareFetch.tryComplete());

Review comment (bug):
For safety, should we be mocking `ShareFetchUtils.processFetchResponse` for the test `testRemoteStorageFetchRequestCompletionOnFutureCompletionFailure`

Code after:
        // Remote fetch related mocks. Remote fetch object completes within tryComplete in this mock, hence request will move on to forceComplete.
        RemoteLogReadResult remoteFetchResult = new RemoteLogReadResult(
            Optional.empty(),
            Optional.of(new TimeoutException("Error occurred while creating remote fetch result")) // Remote fetch result is returned with an error.
        );
        RemoteLogManager remoteLogManager = mock(RemoteLogManager.class);
        doAnswer(invocationOnMock -> {
            // Make sure that the callback is called to populate remoteFetchResult for the mock behaviour.
            Consumer<RemoteLogReadResult> callback = invocationOnMock.getArgument(1);
            callback.accept(remoteFetchResult);
            return CompletableFuture.completedFuture(remoteFetchResult);
        }).when(remoteLogManager).asyncRead(any(), any());
        when(replicaManager.remoteLogManager()).thenReturn(Option.apply(remoteLogManager));

        Uuid fetchId = Uuid.randomUuid();

        Partition p0 = mock(Partition.class);
        when(p0.isLeader()).thenReturn(true);

        Partition p1 = mock(Partition.class);
        when(p1.isLeader()).thenReturn(true);

        when(replicaManager.getPartitionOrException(tp0.topicPartition())).thenReturn(p0);
        when(replicaManager.getPartitionOrException(tp1.topicPartition())).thenReturn(p1);

        BiConsumer<SharePartitionKey, Throwable> exceptionHandler = mockExceptionHandler();
        DelayedShareFetch delayedShareFetch = spy(DelayedShareFetchBuilder.builder()
            .withShareFetchData(shareFetch)
            .withSharePartitions(sharePartitions)
            .withReplicaManager(replicaManager)
            .withPartitionMaxBytesStrategy(mockPartitionMaxBytes(Set.of(tp0, tp1)))
            .withFetchId(fetchId)
            .withExceptionHandler(exceptionHandler)
            .build());

        // sp0 is acquirable, sp1 is not acquirable.
        when(sp0.maybeAcquireFetchLock(fetchId)).thenReturn(true);
        when(sp1.maybeAcquireFetchLock(fetchId)).thenReturn(false);

        try (MockedStatic<ShareFetchUtils> mockedShareFetchUtils = Mockito.mockStatic(ShareFetchUtils.class, Mockito.CALLS_REAL_METHODS)) {
            mockedShareFetchUtils.when(() -> ShareFetchUtils.processFetchResponse(any(), any(), any(), any(), any()))
                .thenReturn(Map.of(tp0, new ShareFetchResponseData.PartitionData().setErrorCode(Errors.REQUEST_TIMED_OUT.code())));

            assertFalse(delayedShareFetch.isCompleted());
            assertTrue(delayedShareFetch.tryComplete());

            assertTrue(delayedShareFetch.isCompleted());
            // Pending remote fetch object gets created for delayed share fetch.
            assertNotNull(delayedShareFetch.pendingRemoteFetches());
            // Verify the locks are released for tp0.
            Mockito.verify(delayedShareFetch, times(1)).releasePartitionLocks(Set.of(tp0));