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

Code before:
        fetchAcquiredRecords(
            sharePartition.acquire(MEMBER_ID, BATCH_SIZE, MAX_FETCH_RECORDS, 0,
                fetchPartitionData(memoryRecords(10, 0)), FETCH_ISOLATION_HWM
            ), 10
        );

        // Acquire a single batch 10-19 with member-1.
        fetchAcquiredRecords(
            sharePartition.acquire(MEMBER_ID, BATCH_SIZE, MAX_FETCH_RECORDS, 10,
                fetchPartitionData(memoryRecords(10, 10)), FETCH_ISOLATION_HWM
            ), 10
        );

        // Validate that there is no ongoing transition.
        assertEquals(2, sharePartition.cachedState().size());
        assertFalse(sharePartition.cachedState().get(0L).batchHasOngoingStateTransition());
        assertFalse(sharePartition.cachedState().get(10L).batchHasOngoingStateTransition());
        assertEquals(RecordState.ACQUIRED, sharePartition.cachedState().get(0L).batchState());
        assertEquals(RecordState.ACQUIRED, sharePartition.cachedState().get(10L).batchState());

        // Return futures which will be completed later, so the batch state has ongoing transition.
        CompletableFuture<WriteShareGroupStateResult> future1 = new CompletableFuture<>();
        CompletableFuture<WriteShareGroupStateResult> future2 = new CompletableFuture<>();

        // Mocking the persister write state RPC to return future 1 when acknowledgement occurs for offsets 0-9.
        // Mocking the persister write state RPC to return future 2 when acknowledgement occurs for offsets 10-19.
        Mockito.when(persister.writeState(Mockito.any())).thenReturn(future1).thenReturn(future2);

        // Acknowledge batch to create ongoing transition.
        sharePartition.acknowledge(MEMBER_ID, List.of(new ShareAcknowledgementBatch(0, 9, List.of(AcknowledgeType.RELEASE.id))));
        sharePartition.acknowledge(MEMBER_ID, List.of(new ShareAcknowledgementBatch(10, 19, List.of(AcknowledgeType.RELEASE.id))));

        // Complete future2 so second acknowledge API can be completed, which updates the cache.
        WriteShareGroupStateResult writeShareGroupStateResult = Mockito.mock(WriteShareGroupStateResult.class);
        Mockito.when(writeShareGroupStateResult.topicsData()).thenReturn(List.of(
            new TopicData<>(TOPIC_ID_PARTITION.topicId(), List.of(
                PartitionFactory.newPartitionErrorData(0, Errors.NONE.code(), Errors.NONE.message())))));
        future2.complete(writeShareGroupStateResult);

        // Offsets 0-9 will have ongoing state transition since future1 is not complete yet.
        // Offsets 10-19 won't have ongoing state transition since future2 has been completed.
        assertTrue(sharePartition.cachedState().get(0L).batchHasOngoingStateTransition());
        assertFalse(sharePartition.cachedState().get(10L).batchHasOngoingStateTransition());

        // nextFetchOffset should return 10 and not 0 because batch 0-9 is undergoing state transition.
        assertEquals(10, sharePartition.nextFetchOffset());
    }

    @Test
    public void testNextFetchOffsetWhenOffsetsHaveOngoingTransition() {
        Persister persister = Mockito.mock(Persister.class);

Review comment (suggestion):
```suggestion
        // Mocking the persister write state RPC to return future 1 and future 2 when acknowledgement occurs for offsets 0-9 and 10-19 respectively.
```

Code after:
        fetchAcquiredRecords(
            sharePartition.acquire(MEMBER_ID, BATCH_SIZE, MAX_FETCH_RECORDS, 0,
                fetchPartitionData(memoryRecords(10, 0)), FETCH_ISOLATION_HWM
            ), 10
        );

        // Acquire a single batch 10-19 with member-1.
        fetchAcquiredRecords(
            sharePartition.acquire(MEMBER_ID, BATCH_SIZE, MAX_FETCH_RECORDS, 10,
                fetchPartitionData(memoryRecords(10, 10)), FETCH_ISOLATION_HWM
            ), 10
        );

        // Validate that there is no ongoing transition.
        assertEquals(2, sharePartition.cachedState().size());
        assertFalse(sharePartition.cachedState().get(0L).batchHasOngoingStateTransition());
        assertFalse(sharePartition.cachedState().get(10L).batchHasOngoingStateTransition());
        assertEquals(RecordState.ACQUIRED, sharePartition.cachedState().get(0L).batchState());
        assertEquals(RecordState.ACQUIRED, sharePartition.cachedState().get(10L).batchState());

        // Return futures which will be completed later, so the batch state has ongoing transition.
        CompletableFuture<WriteShareGroupStateResult> future1 = new CompletableFuture<>();
        CompletableFuture<WriteShareGroupStateResult> future2 = new CompletableFuture<>();

        // Mocking the persister write state RPC to return future 1 and future 2 when acknowledgement occurs for
        // offsets 0-9 and 10-19 respectively.
        Mockito.when(persister.writeState(Mockito.any())).thenReturn(future1).thenReturn(future2);

        // Acknowledge batch to create ongoing transition.
        sharePartition.acknowledge(MEMBER_ID, List.of(new ShareAcknowledgementBatch(0, 9, List.of(AcknowledgeType.RELEASE.id))));
        sharePartition.acknowledge(MEMBER_ID, List.of(new ShareAcknowledgementBatch(10, 19, List.of(AcknowledgeType.RELEASE.id))));

        // Complete future2 so second acknowledge API can be completed, which updates the cache.
        WriteShareGroupStateResult writeShareGroupStateResult = Mockito.mock(WriteShareGroupStateResult.class);
        Mockito.when(writeShareGroupStateResult.topicsData()).thenReturn(List.of(
            new TopicData<>(TOPIC_ID_PARTITION.topicId(), List.of(
                PartitionFactory.newPartitionErrorData(0, Errors.NONE.code(), Errors.NONE.message())))));
        future2.complete(writeShareGroupStateResult);

        // Offsets 0-9 will have ongoing state transition since future1 is not complete yet.
        // Offsets 10-19 won't have ongoing state transition since future2 has been completed.
        assertTrue(sharePartition.cachedState().get(0L).batchHasOngoingStateTransition());
        assertFalse(sharePartition.cachedState().get(10L).batchHasOngoingStateTransition());

        // nextFetchOffset should return 10 and not 0 because batch 0-9 is undergoing state transition.
        assertEquals(10, sharePartition.nextFetchOffset());
    }

    @Test
    public void testNextFetchOffsetWhenOffsetsHaveOngoingTransition() {
        Persister persister = Mockito.mock(Persister.class);