@Test
    public void testWithLeaderChangesAcrossRetries() {
        ProducerBatch batch = new ProducerBatch(new TopicPartition("topic", 1), memoryRecordsBuilder, now);

        // Starting state for the batch, no attempt made to send it yet.
        assertEquals(OptionalInt.empty(), batch.currentLeaderEpoch());
        assertEquals(0, batch.attemptsWhenLeaderLastChanged()); // default value
        batch.maybeUpdateLeaderEpoch(OptionalInt.empty());
        assertFalse(batch.hasLeaderChangedForTheOngoingRetry());

        // 1st attempt[Not a retry] to send the batch.
        // Check leader isn't flagged as a new leader.
        int batchLeaderEpoch = 100;
        batch.maybeUpdateLeaderEpoch(OptionalInt.of(batchLeaderEpoch));
        assertFalse(batch.hasLeaderChangedForTheOngoingRetry(), "batch leader is assigned for 1st time");
        assertEquals(batchLeaderEpoch, batch.currentLeaderEpoch().getAsInt());
        assertEquals(0, batch.attemptsWhenLeaderLastChanged());

        // 2nd attempt[1st retry] to send the batch to a new leader.
        // Check leader change is detected.
        batchLeaderEpoch = 101;
        batch.reenqueued(0);
        batch.maybeUpdateLeaderEpoch(OptionalInt.of(batchLeaderEpoch));
        assertTrue(batch.hasLeaderChangedForTheOngoingRetry(), "batch leader has changed");
        assertEquals(batchLeaderEpoch, batch.currentLeaderEpoch().getAsInt());
        assertEquals(1, batch.attemptsWhenLeaderLastChanged());

        // 2nd attempt[1st retry] still ongoing, yet to be made.
        // Check same leaderEpoch(101) is still considered as a leader-change.
        batch.maybeUpdateLeaderEpoch(OptionalInt.of(batchLeaderEpoch));
        assertTrue(batch.hasLeaderChangedForTheOngoingRetry(), "batch leader has changed");
        assertEquals(batchLeaderEpoch, batch.currentLeaderEpoch().getAsInt());
        assertEquals(1, batch.attemptsWhenLeaderLastChanged());

        // 3rd attempt[2nd retry] to the same leader-epoch(101).
        // Check same leaderEpoch(101) as not detected as a leader-change.
        batch.reenqueued(0);
        batch.maybeUpdateLeaderEpoch(OptionalInt.of(batchLeaderEpoch));
        assertFalse(batch.hasLeaderChangedForTheOngoingRetry(), "batch leader has not changed");
        assertEquals(batchLeaderEpoch, batch.currentLeaderEpoch().getAsInt());
        assertEquals(1, batch.attemptsWhenLeaderLastChanged());

        // Attempt made to update batch leader-epoch to an older leader-epoch(100).
        // Check batch leader-epoch remains unchanged as 101.
        batch.maybeUpdateLeaderEpoch(OptionalInt.of(batchLeaderEpoch - 1));
        assertFalse(batch.hasLeaderChangedForTheOngoingRetry(), "batch leader has not changed");
        assertEquals(batchLeaderEpoch, batch.currentLeaderEpoch().getAsInt());
        assertEquals(1, batch.attemptsWhenLeaderLastChanged());

        // Attempt made to update batch leader-epoch to an unknown leader(optional.empty())
        // Check batch leader-epoch remains unchanged as 101.
        batch.maybeUpdateLeaderEpoch(OptionalInt.empty());
        assertFalse(batch.hasLeaderChangedForTheOngoingRetry(), "batch leader has not changed");
        assertEquals(batchLeaderEpoch, batch.currentLeaderEpoch().getAsInt());
        assertEquals(1, batch.attemptsWhenLeaderLastChanged());
    }
This tests that leader is correctly maintained & leader-change is correctly detected across retries
of the batch. It does so by testing primarily testing methods
1. maybeUpdateLeaderEpoch
2. hasLeaderChangedForTheOngoingRetry
clients/src/test/java/org/apache/kafka/clients/producer/internals/ProducerBatchTest.java