Repository: apache/kafka
File: clients/src/test/java/org/apache/kafka/clients/consumer/internals/ShareConsumeRequestManagerTest.java (Java)

Code before:
        // Validate metadata update is requested due to the leadership error
        assertTrue(metadata.updateRequested());

        // Move the leadership of tp1 onto node 1
        LinkedList<Node> leaderNodes = new LinkedList<>(Arrays.asList(tp0Leader, tp1Leader));
        metadata.updatePartitionLeadership(Map.of(tp1, new Metadata.LeaderIdAndEpoch(Optional.of(nodeId0.id()), Optional.of(validLeaderEpoch + 1))), leaderNodes);

        assertNotEquals(startingClusterMetadata, metadata.fetch());

        // And now the partitions are on the same leader so only one fetch is sent
        assertEquals(1, sendFetches());
        assertFalse(shareConsumeRequestManager.hasCompletedFetches());

        partitionData = buildPartitionDataMap(tip0, records, ShareCompletedFetchTest.acquiredRecords(2L, 1), Errors.NONE, Errors.NONE);
        partitionData.put(tip1,
            new ShareFetchResponseData.PartitionData()
                .setPartitionIndex(tip1.topicPartition().partition())
                .setRecords(records)
                .setAcquiredRecords(ShareCompletedFetchTest.acquiredRecords(1L, 1))
                .setAcknowledgeErrorCode(Errors.NONE.code()));
        client.prepareResponseFrom(ShareFetchResponse.of(Errors.NONE, 0, partitionData, Collections.emptyList(), 0), nodeId0);
        networkClientDelegate.poll(time.timer(0));
        assertTrue(shareConsumeRequestManager.hasCompletedFetches());

        partitionRecords = fetchRecords();
        assertTrue(partitionRecords.containsKey(tp0) && partitionRecords.containsKey(tp1));
        assertEquals(1, partitionRecords.get(tp0).size());
        assertEquals(1, partitionRecords.get(tp1).size());
    }

    /**
     * Test the scenario that ShareFetchResponse returns with an error indicating leadership change for the partition,
     * along with new leader info (defined in KIP-951).
     */
    @ParameterizedTest
    @EnumSource(value = Errors.class, names = {"FENCED_LEADER_EPOCH", "NOT_LEADER_OR_FOLLOWER"})
    public void testWhenFetchResponseReturnsWithALeadershipChangeErrorAndNewLeaderInformation(Errors error) {
        buildRequestManager();

        subscriptions.subscribeToShareGroup(Collections.singleton(topicName));
        Set<TopicPartition> partitions = new HashSet<>();
        partitions.add(tp0);
        partitions.add(tp1);
        subscriptions.assignFromSubscribed(partitions);

        client.updateMetadata(
            RequestTestUtils.metadataUpdateWithIds(2, Map.of(topicName, 2),
                tp -> validLeaderEpoch, topicIds, false));
        Node nodeId0 = metadata.fetch().nodeById(0);
        Node nodeId1 = metadata.fetch().nodeById(1);
        Node tp0Leader = metadata.fetch().leaderFor(tp0);

Review comment (suggestion):
I would leave this line unchanged as two assertions just to minimise the amount of change.

Code after:
        List<ConsumerRecord<byte[], byte[]>> fetchedRecords = partitionRecords.get(tp0);
        assertEquals(1, fetchedRecords.size());

        Acknowledgements acknowledgements = Acknowledgements.empty();
        acknowledgements.add(1L, AcknowledgeType.ACCEPT);
        shareConsumeRequestManager.fetch(Map.of(tip0, new NodeAcknowledgements(0, acknowledgements)));

        assertEquals(startingClusterMetadata, metadata.fetch());

        // Validate metadata update is requested due to the leadership error
        assertTrue(metadata.updateRequested());

        // Move the leadership of tp1 onto node 1
        LinkedList<Node> leaderNodes = new LinkedList<>(Arrays.asList(tp0Leader, tp1Leader));
        metadata.updatePartitionLeadership(Map.of(tp1, new Metadata.LeaderIdAndEpoch(Optional.of(nodeId0.id()), Optional.of(validLeaderEpoch + 1))), leaderNodes);

        assertNotEquals(startingClusterMetadata, metadata.fetch());

        // And now the partitions are on the same leader so only one fetch is sent
        assertEquals(1, sendFetches());
        assertFalse(shareConsumeRequestManager.hasCompletedFetches());

        partitionData = buildPartitionDataMap(tip0, records, ShareCompletedFetchTest.acquiredRecords(2L, 1), Errors.NONE, Errors.NONE);
        partitionData.put(tip1,
            new ShareFetchResponseData.PartitionData()
                .setPartitionIndex(tip1.topicPartition().partition())
                .setRecords(records)
                .setAcquiredRecords(ShareCompletedFetchTest.acquiredRecords(1L, 1))
                .setAcknowledgeErrorCode(Errors.NONE.code()));
        client.prepareResponseFrom(ShareFetchResponse.of(Errors.NONE, 0, partitionData, Collections.emptyList(), 0), nodeId0);
        networkClientDelegate.poll(time.timer(0));
        assertTrue(shareConsumeRequestManager.hasCompletedFetches());

        partitionRecords = fetchRecords();
        assertTrue(partitionRecords.containsKey(tp0));
        assertTrue(partitionRecords.containsKey(tp1));

        fetchedRecords = partitionRecords.get(tp0);
        assertEquals(1, fetchedRecords.size());
        fetchedRecords = partitionRecords.get(tp1);
        assertEquals(1, fetchedRecords.size());
    }

    /**
     * Test the scenario that ShareFetchResponse returns with an error indicating leadership change for the partition,
     * along with new leader info (defined in KIP-951).
     */
    @ParameterizedTest
    @EnumSource(value = Errors.class, names = {"FENCED_LEADER_EPOCH", "NOT_LEADER_OR_FOLLOWER"})
    public void testWhenFetchResponseReturnsWithALeadershipChangeErrorAndNewLeaderInformation(Errors error) {
        buildRequestManager();