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

Code before:
    }

    @ParameterizedTest
    @ValueSource(booleans = { true, false })
    public void testUnattachedWithLeaderCanBecomeFollowerAfterFindingLeader(boolean withKip853Rpc) throws Exception {
        int localId = randomReplicaId();
        int otherNodeId = localId + 1;
        int leaderNodeId = localId + 2;
        int epoch = 5;
        Set<Integer> voters = Set.of(localId, otherNodeId, leaderNodeId);
        List<InetSocketAddress> bootstrapServers = voters
            .stream()
            .map(RaftClientTestContext::mockAddress)
            .toList();

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

        // after fetch timeout, node will become prospective with leader
        context.time.sleep(context.fetchTimeoutMs);
        context.client.poll();
        assertTrue(context.client.quorum().isProspective());
        assertEquals(leaderNodeId, context.client.quorum().leaderId().getAsInt());
        assertEquals(epoch, context.currentEpoch());

        // after election loss node will become follower because it had a last known leader
        context.time.sleep(context.electionTimeoutMs() * 2L);
        context.client.poll();
        assertTrue(context.client.quorum().isFollower());
        assertEquals(leaderNodeId, context.client.quorum().leaderId().getAsInt());
        assertEquals(epoch, context.currentEpoch());

        // node will send fetch request to leader
        context.pollUntilRequest();
        RaftRequest.Outbound request = context.assertSentFetchRequest(epoch, 0L, 0);
        assertTrue(context.client.quorum().isFollower());
        assertTrue(context.client.quorum().isVoter());
        assertEquals(leaderNodeId, request.destination().id());
    }

    @ParameterizedTest
    @ValueSource(booleans = { true, false })
    public void testInitializeObserverNoPreviousState(boolean withKip853Rpc) throws Exception {
        int localId = randomReplicaId();
        int leaderId = localId + 1;
        int otherNodeId = localId + 2;
        int epoch = 5;
        Set<Integer> voters = Set.of(leaderId, otherNodeId);

Review comment (suggestion):
You can use `RaftClientTestContext#assertElectedLeader(int, int)`.

Code after:
    }

    @ParameterizedTest
    @ValueSource(booleans = { true, false })
    public void testInitializeObserverNoPreviousState(boolean withKip853Rpc) throws Exception {
        int localId = randomReplicaId();
        int leaderId = localId + 1;
        int otherNodeId = localId + 2;
        int epoch = 5;
        Set<Integer> voters = Set.of(leaderId, otherNodeId);

        RaftClientTestContext context = new RaftClientTestContext.Builder(localId, voters)
            .withKip853Rpc(withKip853Rpc)
            .build();

        context.pollUntilRequest();
        RaftRequest.Outbound fetchRequest = context.assertSentFetchRequest();
        assertTrue(voters.contains(fetchRequest.destination().id()));
        context.assertFetchRequestData(fetchRequest, 0, 0L, 0, context.client.highWatermark());

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

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

    @ParameterizedTest
    @ValueSource(booleans = { true, false })
    public void testObserverQuorumDiscoveryFailure(boolean withKip853Rpc) throws Exception {
        int localId = randomReplicaId();
        int leaderId = localId + 1;
        int epoch = 5;
        Set<Integer> voters = Set.of(leaderId);
        List<InetSocketAddress> bootstrapServers = voters
            .stream()
            .map(RaftClientTestContext::mockAddress)
            .toList();

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

        context.pollUntilRequest();
        RaftRequest.Outbound fetchRequest = context.assertSentFetchRequest();
        assertTrue(context.bootstrapIds.contains(fetchRequest.destination().id()));
        context.assertFetchRequestData(fetchRequest, 0, 0L, 0, context.client.highWatermark());