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

Code before:
        when(pendingRemoteFetches.remoteFetches()).thenReturn(List.of(remoteFetch));
        when(pendingRemoteFetches.isDone()).thenReturn(false);

        // Make sure that the callback to complete remote storage share fetch result is not called.
        doAnswer(invocationOnMock -> null).when(pendingRemoteFetches).invokeCallbackOnCompletion(any());

        // Mock the behaviour of replica manager such that remote storage fetch completion timer task completes on adding it to the watch queue.
        doAnswer(invocationOnMock -> {
            TimerTask timerTask = invocationOnMock.getArgument(0);
            timerTask.run();
            return null;
        }).when(replicaManager).addShareFetchTimerRequest(any());

        try (MockedStatic<ShareFetchUtils> mockedShareFetchUtils = Mockito.mockStatic(ShareFetchUtils.class)) {
            Map<TopicIdPartition, ShareFetchResponseData.PartitionData> partitionDataMap = new LinkedHashMap<>();
            partitionDataMap.put(tp0, mock(ShareFetchResponseData.PartitionData.class));
            mockedShareFetchUtils.when(() -> ShareFetchUtils.processFetchResponse(any(), any(), any(), any(), any())).thenReturn(partitionDataMap);

            assertFalse(delayedShareFetch.isCompleted());
            delayedShareFetch.forceComplete();
            assertTrue(delayedShareFetch.isCompleted());
            // the future of shareFetch completes.
            assertTrue(shareFetch.isCompleted());
            assertEquals(Set.of(tp0), future.join().keySet());
            // Verify the locks are released for tp0.
            Mockito.verify(delayedShareFetch, times(1)).releasePartitionLocks(Set.of(tp0));
            assertTrue(delayedShareFetch.outsidePurgatoryCallbackLock());
            assertTrue(delayedShareFetch.lock().tryLock());
            delayedShareFetch.lock().unlock();
        }
    }

    static void mockTopicIdPartitionToReturnDataEqualToMinBytes(ReplicaManager replicaManager, TopicIdPartition topicIdPartition, int minBytes) {
        LogOffsetMetadata hwmOffsetMetadata = new LogOffsetMetadata(1, 1, minBytes);
        LogOffsetSnapshot endOffsetSnapshot = new LogOffsetSnapshot(1, mock(LogOffsetMetadata.class),
            hwmOffsetMetadata, mock(LogOffsetMetadata.class));
        Partition partition = mock(Partition.class);
        when(partition.isLeader()).thenReturn(true);
        when(partition.getLeaderEpoch()).thenReturn(1);
        when(partition.fetchOffsetSnapshot(any(), anyBoolean())).thenReturn(endOffsetSnapshot);
        when(replicaManager.getPartitionOrException(topicIdPartition.topicPartition())).thenReturn(partition);
    }

    private PartitionMaxBytesStrategy mockPartitionMaxBytes(Set<TopicIdPartition> partitions) {
        PartitionMaxBytesStrategy partitionMaxBytesStrategy = mock(PartitionMaxBytesStrategy.class);
        LinkedHashMap<TopicIdPartition, Integer> maxBytes = new LinkedHashMap<>();
        partitions.forEach(partition -> maxBytes.put(partition, 1));
        when(partitionMaxBytesStrategy.maxBytes(anyInt(), any(), anyInt())).thenReturn(maxBytes);
        return partitionMaxBytesStrategy;
    }


Review comment (suggestion):
Removal of this function doesn't make sense to me. I agree that the line `when(replicaManager.getPartitionOrException(topicIdPartition.topicPartition())).thenReturn(partition);` is problematic. We can pass the mock Partition object into this function. Example - 

caller function - 
```
Partition p0 = mock(Partition.class);
mockTopicIdPartitionFetchBytes(replicaManager, tp0, hwmOffsetMetadata, p0);
when(p0.isLeader()).thenReturn(true);
when(replicaManager.getPartitionOrException(tp0.topicPartition())).thenReturn(p0);
```

function `mockTopicIdPartitionFetchBytes` - 

```
private void mockTopicIdPartitionFetchBytes(LogOffsetMetadata hwmOffsetMetadata, Partition partition) {
        LogOffsetSnapshot endOffsetSnapshot = new LogOffsetSnapshot(1, mock(LogOffsetMetadata.class),
            hwmOffsetMetadata, mock(LogOffsetMetadata.class));
        when(partition.fetchOffsetSnapshot(any(), anyBoolean())).thenReturn(endOffsetSnapshot);
}
```


Code after:

        // Mock the behaviour of replica manager such that remote storage fetch completion timer task completes on adding it to the watch queue.
        doAnswer(invocationOnMock -> {
            TimerTask timerTask = invocationOnMock.getArgument(0);
            timerTask.run();
            return null;
        }).when(replicaManager).addShareFetchTimerRequest(any());

        try (MockedStatic<ShareFetchUtils> mockedShareFetchUtils = Mockito.mockStatic(ShareFetchUtils.class)) {
            Map<TopicIdPartition, ShareFetchResponseData.PartitionData> partitionDataMap = new LinkedHashMap<>();
            partitionDataMap.put(tp0, mock(ShareFetchResponseData.PartitionData.class));
            mockedShareFetchUtils.when(() -> ShareFetchUtils.processFetchResponse(any(), any(), any(), any(), any())).thenReturn(partitionDataMap);

            assertFalse(delayedShareFetch.isCompleted());
            delayedShareFetch.forceComplete();
            assertTrue(delayedShareFetch.isCompleted());
            // the future of shareFetch completes.
            assertTrue(shareFetch.isCompleted());
            assertEquals(Set.of(tp0), future.join().keySet());
            // Verify the locks are released for tp0.
            Mockito.verify(delayedShareFetch, times(1)).releasePartitionLocks(Set.of(tp0));
            assertTrue(delayedShareFetch.outsidePurgatoryCallbackLock());
            assertTrue(delayedShareFetch.lock().tryLock());
            delayedShareFetch.lock().unlock();
        }
    }

    static void mockTopicIdPartitionToReturnDataEqualToMinBytes(ReplicaManager replicaManager, TopicIdPartition topicIdPartition, int minBytes) {
        LogOffsetMetadata hwmOffsetMetadata = new LogOffsetMetadata(1, 1, minBytes);
        LogOffsetSnapshot endOffsetSnapshot = new LogOffsetSnapshot(1, mock(LogOffsetMetadata.class),
            hwmOffsetMetadata, mock(LogOffsetMetadata.class));
        Partition partition = mock(Partition.class);
        when(partition.isLeader()).thenReturn(true);
        when(partition.getLeaderEpoch()).thenReturn(1);
        when(partition.fetchOffsetSnapshot(any(), anyBoolean())).thenReturn(endOffsetSnapshot);
        when(replicaManager.getPartitionOrException(topicIdPartition.topicPartition())).thenReturn(partition);
    }

    private void mockTopicIdPartitionFetchBytes(LogOffsetMetadata hwmOffsetMetadata, Partition partition) {
        LogOffsetSnapshot endOffsetSnapshot = new LogOffsetSnapshot(1, mock(LogOffsetMetadata.class),
            hwmOffsetMetadata, mock(LogOffsetMetadata.class));
        when(partition.fetchOffsetSnapshot(any(), anyBoolean())).thenReturn(endOffsetSnapshot);
    }

    private PartitionMaxBytesStrategy mockPartitionMaxBytes(Set<TopicIdPartition> partitions) {
        PartitionMaxBytesStrategy partitionMaxBytesStrategy = mock(PartitionMaxBytesStrategy.class);
        LinkedHashMap<TopicIdPartition, Integer> maxBytes = new LinkedHashMap<>();
        partitions.forEach(partition -> maxBytes.put(partition, 1));
        when(partitionMaxBytesStrategy.maxBytes(anyInt(), any(), anyInt())).thenReturn(maxBytes);
        return partitionMaxBytesStrategy;
    }