Repository: apache/kafka
File: raft/src/test/java/org/apache/kafka/raft/KafkaRaftClientTest.java (Java)

Code before:
            context.fetchResponse(epoch, leaderId, MemoryRecords.EMPTY, 0L, Errors.FENCED_LEADER_EPOCH)
        );

        context.client.poll();
        context.assertElectedLeader(epoch, leaderId);

        context.time.sleep(context.fetchTimeoutMs);

        context.pollUntilRequest();
        fetchRequest = context.assertSentFetchRequest();
        assertNotEquals(leaderId, fetchRequest.destination().id());
        assertTrue(context.bootstrapIds.contains(fetchRequest.destination().id()));
        assertTrue(context.client.quorum().followerStateOrThrow().hasFetchTimeoutExpired(context.time.milliseconds()));
        context.assertFetchRequestData(fetchRequest, epoch, 0L, 0);

        context.deliverResponse(
            fetchRequest.correlationId(),
            fetchRequest.destination(),
            context.fetchResponse(epoch, leaderId, MemoryRecords.EMPTY, 0L, Errors.NOT_LEADER_OR_FOLLOWER)
        );

        context.client.poll();
        context.pollUntilRequest();
        fetchRequest = context.assertSentFetchRequest();
        assertEquals(leaderId, fetchRequest.destination().id());
        assertFalse(context.client.quorum().followerStateOrThrow().hasFetchTimeoutExpired(context.time.milliseconds()));
    }

    @ParameterizedTest
    @ValueSource(booleans = { true, false })
    public void testObserverHandleRetryFetchToBootstrapServer(boolean withKip853Rpc) throws Exception {
        // This test tries to check that KRaft is able to handle a retrying Fetch request to
        // a boostrap server after a Fetch request to the leader.
        int localId = randomReplicaId();
        int leaderId = localId + 1;
        int otherNodeId = localId + 2;
        int epoch = 5;
        Set<Integer> voters = Set.of(leaderId, otherNodeId);
        List<InetSocketAddress> bootstrapServers = voters
            .stream()
            .map(RaftClientTestContext::mockAddress)
            .toList();

        RaftClientTestContext context = new RaftClientTestContext.Builder(localId, voters)
            .withBootstrapServers(Optional.of(bootstrapServers))
            .withKip853Rpc(withKip853Rpc)
            .build();

        // Expect a fetch request to one of the bootstrap servers
        context.pollUntilRequest();
        RaftRequest.Outbound discoveryFetchRequest = context.assertSentFetchRequest();

Review comment (refactor):
Let's remove this line. This suite should test the protocol and not the internal implementation state.

Code after:
        assertTrue(context.bootstrapIds.contains(fetchRequest.destination().id()));
        context.assertFetchRequestData(fetchRequest, epoch, 0L, 0, context.client.highWatermark());

        // observer will retry to the leader
        context.time.sleep(context.requestTimeoutMs());
        context.client.poll();
        fetchRequest = context.assertSentFetchRequest();
        assertEquals(leaderId, fetchRequest.destination().id());
        context.assertFetchRequestData(fetchRequest, epoch, 0L, 0, context.client.highWatermark());
    }

    @ParameterizedTest
    @ValueSource(booleans = { true, false })
    public void testObserverHandleRetryFetchToBootstrapServer(boolean withKip853Rpc) throws Exception {
        // This test tries to check that KRaft is able to handle a retrying Fetch request to
        // a boostrap server after a Fetch request to the leader.
        int localId = randomReplicaId();
        int leaderId = localId + 1;
        int otherNodeId = localId + 2;
        int epoch = 5;
        Set<Integer> voters = Set.of(leaderId, otherNodeId);
        List<InetSocketAddress> bootstrapServers = voters
            .stream()
            .map(RaftClientTestContext::mockAddress)
            .toList();

        RaftClientTestContext context = new RaftClientTestContext.Builder(localId, voters)
            .withBootstrapServers(Optional.of(bootstrapServers))
            .withKip853Rpc(withKip853Rpc)
            .build();

        // Expect a fetch request to one of the bootstrap servers
        context.pollUntilRequest();
        RaftRequest.Outbound discoveryFetchRequest = context.assertSentFetchRequest();
        assertFalse(voters.contains(discoveryFetchRequest.destination().id()));
        assertTrue(context.bootstrapIds.contains(discoveryFetchRequest.destination().id()));
        context.assertFetchRequestData(discoveryFetchRequest, 0, 0L, 0, context.client.highWatermark());

        // Send a response with the leader and epoch
        context.deliverResponse(
            discoveryFetchRequest.correlationId(),
            discoveryFetchRequest.destination(),
            context.fetchResponse(epoch, leaderId, MemoryRecords.EMPTY, 0L, Errors.FENCED_LEADER_EPOCH)
        );

        context.client.poll();
        context.assertElectedLeader(epoch, leaderId);

        // Expect a fetch request to the leader
        context.pollUntilRequest();
        RaftRequest.Outbound toLeaderFetchRequest = context.assertSentFetchRequest();