/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.kafka.raft.internals;

import org.apache.kafka.common.compress.Compression;
import org.apache.kafka.common.memory.MemoryPool;
import org.apache.kafka.common.message.KRaftVersionRecord;
import org.apache.kafka.common.message.LeaderChangeMessage;
import org.apache.kafka.common.message.SnapshotHeaderRecord;
import org.apache.kafka.common.protocol.ObjectSerializationCache;
import org.apache.kafka.common.protocol.Writable;
import org.apache.kafka.common.record.ControlRecordUtils;
import org.apache.kafka.common.record.DefaultRecord;
import org.apache.kafka.common.record.MemoryRecordsBuilder;
import org.apache.kafka.common.record.RecordBatch;
import org.apache.kafka.common.record.TimestampType;
import org.apache.kafka.common.utils.ByteUtils;
import org.apache.kafka.common.utils.MockTime;
import org.apache.kafka.common.utils.Utils;

import org.junit.jupiter.api.Test;
import org.mockito.Mockito;

import java.nio.ByteBuffer;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;

import static java.util.Arrays.asList;
import static java.util.Collections.singletonList;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotEquals;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;

class BatchAccumulatorTest {
    private final MemoryPool memoryPool = Mockito.mock(MemoryPool.class);
    private final MockTime time = new MockTime();
    private final StringSerde serde = new StringSerde();
    private final int maxNumberOfBatches = 10;

    private BatchAccumulator<String> buildAccumulator(
        int leaderEpoch,
        long baseOffset,
        int lingerMs,
        int maxBatchSize
    ) {
        return new BatchAccumulator<>(
            leaderEpoch,
            baseOffset,
            lingerMs,
            maxBatchSize,
            maxNumberOfBatches,
            memoryPool,
            time,
            Compression.NONE,
            serde
        );
    }

    @Test
    public void testLeaderChangeMessageWritten() {
        int leaderEpoch = 17;
        long baseOffset = 0;
        int lingerMs = 50;
        int maxBatchSize = 512;

        ByteBuffer buffer = ByteBuffer.allocate(maxBatchSize);
        Mockito.when(memoryPool.tryAllocate(maxBatchSize))
            .thenReturn(buffer);

        BatchAccumulator<String> acc = buildAccumulator(
            leaderEpoch,
            baseOffset,
            lingerMs,
            maxBatchSize
        );

        acc.appendLeaderChangeMessage(new LeaderChangeMessage(), time.milliseconds());
        assertTrue(acc.needsDrain(time.milliseconds()));

        List<BatchAccumulator.CompletedBatch<String>> batches = acc.drain();
        assertEquals(1, batches.size());

        BatchAccumulator.CompletedBatch<String> batch = batches.get(0);
        batch.release();
        Mockito.verify(memoryPool).release(buffer);
    }

    @Test
    public void testForceDrain() {
        int leaderEpoch = 17;
        long baseOffset = 157;
        int lingerMs = 50;
        int maxBatchSize = 512;

        Mockito.when(memoryPool.tryAllocate(maxBatchSize))
            .thenReturn(ByteBuffer.allocate(maxBatchSize));

        BatchAccumulator<String> acc = buildAccumulator(
            leaderEpoch,
            baseOffset,
            lingerMs,
            maxBatchSize
        );

        List<String> records = asList("a", "b", "c", "d", "e", "f", "g", "h", "i");

        // Append records
        assertEquals(baseOffset, acc.append(leaderEpoch, records.subList(0, 1), false));
        assertEquals(baseOffset + 2, acc.append(leaderEpoch, records.subList(1, 3), false));
        assertEquals(baseOffset + 5, acc.append(leaderEpoch, records.subList(3, 6), false));
        assertEquals(baseOffset + 7, acc.append(leaderEpoch, records.subList(6, 8), false));
        assertEquals(baseOffset + 8, acc.append(leaderEpoch, records.subList(8, 9), false));

        assertFalse(acc.needsDrain(time.milliseconds()));
        acc.forceDrain();
        assertTrue(acc.needsDrain(time.milliseconds()));
        assertEquals(0, acc.timeUntilDrain(time.milliseconds()));

        // Drain completed batches
        List<BatchAccumulator.CompletedBatch<String>> batches = acc.drain();

        assertEquals(1, batches.size());
        assertFalse(acc.needsDrain(time.milliseconds()));
        assertEquals(Long.MAX_VALUE - time.milliseconds(), acc.timeUntilDrain(time.milliseconds()));

        BatchAccumulator.CompletedBatch<String> batch = batches.get(0);
        assertEquals(records, batch.records.get());
        assertEquals(baseOffset, batch.baseOffset);
        assertEquals(time.milliseconds(), batch.appendTimestamp());
    }

    @Test
    public void testForceDrainBeforeAppendLeaderChangeMessage() {
        int leaderEpoch = 17;
        long baseOffset = 157;
        int lingerMs = 50;
        int maxBatchSize = 512;

        Mockito.when(memoryPool.tryAllocate(maxBatchSize))
            .thenReturn(ByteBuffer.allocate(maxBatchSize));
        Mockito.when(memoryPool.tryAllocate(256))
            .thenReturn(ByteBuffer.allocate(256));

        BatchAccumulator<String> acc = buildAccumulator(
            leaderEpoch,
            baseOffset,
            lingerMs,
            maxBatchSize
        );

        List<String> records = asList("a", "b", "c", "d", "e", "f", "g", "h", "i");

        // Append records
        assertEquals(baseOffset, acc.append(leaderEpoch, records.subList(0, 1), false));
        assertEquals(baseOffset + 2, acc.append(leaderEpoch, records.subList(1, 3), false));
        assertEquals(baseOffset + 5, acc.append(leaderEpoch, records.subList(3, 6), false));
        assertEquals(baseOffset + 7, acc.append(leaderEpoch, records.subList(6, 8), false));
        assertEquals(baseOffset + 8, acc.append(leaderEpoch, records.subList(8, 9), false));

        assertFalse(acc.needsDrain(time.milliseconds()));

        // Append a leader change message
        acc.appendLeaderChangeMessage(new LeaderChangeMessage(), time.milliseconds());

        assertTrue(acc.needsDrain(time.milliseconds()));

        // Test that drain status is FINISHED
        assertEquals(0, acc.timeUntilDrain(time.milliseconds()));

        // Drain completed batches
        List<BatchAccumulator.CompletedBatch<String>> batches = acc.drain();

        // Should have 2 batches, one consisting of `records` and one `leaderChangeMessage`
        assertEquals(2, batches.size());
        assertFalse(acc.needsDrain(time.milliseconds()));
        assertEquals(Long.MAX_VALUE - time.milliseconds(), acc.timeUntilDrain(time.milliseconds()));

        BatchAccumulator.CompletedBatch<String> batch = batches.get(0);
        assertEquals(records, batch.records.get());
        assertEquals(baseOffset, batch.baseOffset);
        assertEquals(time.milliseconds(), batch.appendTimestamp());
    }

    @Test
    public void testLingerIgnoredIfAccumulatorEmpty() {
        int leaderEpoch = 17;
        long baseOffset = 157;
        int lingerMs = 50;
        int maxBatchSize = 512;

        BatchAccumulator<String> acc = buildAccumulator(
            leaderEpoch,
            baseOffset,
            lingerMs,
            maxBatchSize
        );

        assertTrue(acc.isEmpty());
        assertFalse(acc.needsDrain(time.milliseconds()));
        assertEquals(Long.MAX_VALUE - time.milliseconds(), acc.timeUntilDrain(time.milliseconds()));
    }

    @Test
    public void testLingerBeginsOnFirstWrite() {
        int leaderEpoch = 17;
        long baseOffset = 157;
        int lingerMs = 50;
        int maxBatchSize = 512;

        Mockito.when(memoryPool.tryAllocate(maxBatchSize))
            .thenReturn(ByteBuffer.allocate(maxBatchSize));

        BatchAccumulator<String> acc = buildAccumulator(
            leaderEpoch,
            baseOffset,
            lingerMs,
            maxBatchSize
        );

        time.sleep(15);
        assertEquals(baseOffset, acc.append(leaderEpoch, singletonList("a"), false));
        assertEquals(lingerMs, acc.timeUntilDrain(time.milliseconds()));
        assertFalse(acc.isEmpty());

        time.sleep(lingerMs / 2);
        assertEquals(lingerMs / 2, acc.timeUntilDrain(time.milliseconds()));
        assertFalse(acc.isEmpty());

        time.sleep(lingerMs / 2);
        assertEquals(0, acc.timeUntilDrain(time.milliseconds()));
        assertTrue(acc.needsDrain(time.milliseconds()));
        assertFalse(acc.isEmpty());
    }

    @Test
    public void testCompletedBatchReleaseBuffer() {
        int leaderEpoch = 17;
        long baseOffset = 157;
        int lingerMs = 50;
        int maxBatchSize = 512;

        ByteBuffer buffer = ByteBuffer.allocate(maxBatchSize);
        Mockito.when(memoryPool.tryAllocate(maxBatchSize))
            .thenReturn(buffer);

        BatchAccumulator<String> acc = buildAccumulator(
            leaderEpoch,
            baseOffset,
            lingerMs,
            maxBatchSize
        );

        assertEquals(baseOffset, acc.append(leaderEpoch, singletonList("a"), false));
        time.sleep(lingerMs);

        List<BatchAccumulator.CompletedBatch<String>> batches = acc.drain();
        assertEquals(1, batches.size());

        BatchAccumulator.CompletedBatch<String> batch = batches.get(0);
        batch.release();
        Mockito.verify(memoryPool).release(buffer);
    }

    @Test
    public void testUnflushedBuffersReleasedByClose() {
        int leaderEpoch = 17;
        long baseOffset = 157;
        int lingerMs = 50;
        int maxBatchSize = 512;

        ByteBuffer buffer = ByteBuffer.allocate(maxBatchSize);
        Mockito.when(memoryPool.tryAllocate(maxBatchSize))
            .thenReturn(buffer);

        BatchAccumulator<String> acc = buildAccumulator(
            leaderEpoch,
            baseOffset,
            lingerMs,
            maxBatchSize
        );

        assertEquals(baseOffset, acc.append(leaderEpoch, singletonList("a"), false));
        acc.close();
        Mockito.verify(memoryPool).release(buffer);
    }

    @Test
    public void testSingleBatchAccumulation() {
        int leaderEpoch = 17;
        long baseOffset = 157;
        int lingerMs = 50;
        int maxBatchSize = 512;

        Mockito.when(memoryPool.tryAllocate(maxBatchSize))
            .thenReturn(ByteBuffer.allocate(maxBatchSize));

        BatchAccumulator<String> acc = buildAccumulator(
            leaderEpoch,
            baseOffset,
            lingerMs,
            maxBatchSize
        );

        List<String> records = asList("a", "b", "c", "d", "e", "f", "g", "h", "i");
        assertEquals(baseOffset, acc.append(leaderEpoch, records.subList(0, 1), false));
        assertEquals(baseOffset + 2, acc.append(leaderEpoch, records.subList(1, 3), false));
        assertEquals(baseOffset + 5, acc.append(leaderEpoch, records.subList(3, 6), false));
        assertEquals(baseOffset + 7, acc.append(leaderEpoch, records.subList(6, 8), false));
        assertEquals(baseOffset + 8, acc.append(leaderEpoch, records.subList(8, 9), false));

        long expectedAppendTimestamp = time.milliseconds();
        time.sleep(lingerMs);
        assertTrue(acc.needsDrain(time.milliseconds()));

        List<BatchAccumulator.CompletedBatch<String>> batches = acc.drain();
        assertEquals(1, batches.size());
        assertFalse(acc.needsDrain(time.milliseconds()));
        assertEquals(Long.MAX_VALUE - time.milliseconds(), acc.timeUntilDrain(time.milliseconds()));

        BatchAccumulator.CompletedBatch<String> batch = batches.get(0);
        assertEquals(records, batch.records.get());
        assertEquals(baseOffset, batch.baseOffset);
        assertEquals(expectedAppendTimestamp, batch.appendTimestamp());
    }

    @Test
    public void testMultipleBatchAccumulation() {
        int leaderEpoch = 17;
        long baseOffset = 157;
        int lingerMs = 50;
        int maxBatchSize = 256;
        int numberOfRecords = 100;

        Mockito.when(memoryPool.tryAllocate(maxBatchSize))
            .thenReturn(ByteBuffer.allocate(maxBatchSize));

        BatchAccumulator<String> acc = buildAccumulator(
            leaderEpoch,
            baseOffset,
            lingerMs,
            maxBatchSize
        );

        // Append enough records so that multiple batches get created
        for (int records = 0; records < numberOfRecords; records++) {
            acc.append(leaderEpoch, singletonList("foo"), false);
        }

        List<BatchAccumulator.CompletedBatch<String>> batches = acc.drain();
        assertTrue(batches.size() > 1, () -> String.format("number of batches is %s", batches.size()));
        assertTrue(batches.stream().allMatch(batch -> batch.data.sizeInBytes() <= maxBatchSize));
    }

    @Test
    public void testCloseWhenEmpty() {
        int leaderEpoch = 17;
        long baseOffset = 157;
        int lingerMs = 50;
        int maxBatchSize = 256;

        BatchAccumulator<String> acc = buildAccumulator(
            leaderEpoch,
            baseOffset,
            lingerMs,
            maxBatchSize
        );

        acc.close();
        Mockito.verifyNoInteractions(memoryPool);
    }

    @Test
    public void testDrainDoesNotBlockWithConcurrentAppend() throws Exception {
        int leaderEpoch = 17;
        long baseOffset = 157;
        int lingerMs = 50;
        int maxBatchSize = 256;
        int maxNumberOfBatches = 10;

        StringSerde serde = Mockito.spy(new StringSerde());
        BatchAccumulator<String> acc = new BatchAccumulator<>(
            leaderEpoch,
            baseOffset,
            lingerMs,
            maxBatchSize,
            maxNumberOfBatches,
            memoryPool,
            time,
            Compression.NONE,
            serde
        );

        CountDownLatch acquireLockLatch = new CountDownLatch(1);
        CountDownLatch releaseLockLatch = new CountDownLatch(1);

        // Do the first append outside the thread to start the linger timer
        Mockito.when(memoryPool.tryAllocate(maxBatchSize))
            .thenReturn(ByteBuffer.allocate(maxBatchSize));
        acc.append(leaderEpoch, singletonList("a"), false);

        // Let the serde block to simulate a slow append
        Mockito.doAnswer(invocation -> {
            Writable writable = invocation.getArgument(2);
            acquireLockLatch.countDown();
            releaseLockLatch.await();
            writable.writeByteArray(Utils.utf8("b"));
            return null;
        }).when(serde).write(
            Mockito.eq("b"),
            Mockito.any(ObjectSerializationCache.class),
            Mockito.any(Writable.class)
        );

        Thread appendThread = new Thread(() -> acc.append(leaderEpoch, singletonList("b"), false));
        appendThread.start();

        // Attempt to drain while the append thread is holding the lock
        acquireLockLatch.await();
        time.sleep(lingerMs);
        assertTrue(acc.needsDrain(time.milliseconds()));
        assertEquals(Collections.emptyList(), acc.drain());
        assertTrue(acc.needsDrain(time.milliseconds()));

        // Now let the append thread complete and verify that we can finish the drain
        releaseLockLatch.countDown();
        appendThread.join();
        List<BatchAccumulator.CompletedBatch<String>> drained = acc.drain();
        assertEquals(1, drained.size());
        assertEquals(Long.MAX_VALUE - time.milliseconds(), acc.timeUntilDrain(time.milliseconds()));
        drained.forEach(completedBatch ->
            completedBatch.data.batches().forEach(recordBatch ->
                assertEquals(leaderEpoch, recordBatch.partitionLeaderEpoch()))
        );

        acc.close();
    }

    int recordSizeInBytes(String record, int numberOfRecords) {
        int serdeSize = serde.recordSize(record, new ObjectSerializationCache());

        int recordSizeInBytes = DefaultRecord.sizeOfBodyInBytes(
            numberOfRecords,
            0,
            -1,
            serdeSize,
            DefaultRecord.EMPTY_HEADERS
        );

        return ByteUtils.sizeOfVarint(recordSizeInBytes) + recordSizeInBytes;
    }

    @Test
    public void testMultipleControlRecords() {
        int leaderEpoch = 17;
        long baseOffset = 157;
        int lingerMs = 50;
        int maxBatchSize = 512;

        ByteBuffer buffer = ByteBuffer.allocate(maxBatchSize);
        Mockito.when(memoryPool.tryAllocate(maxBatchSize))
                .thenReturn(buffer);

        try (BatchAccumulator<String> acc = buildAccumulator(
                leaderEpoch,
                baseOffset,
                lingerMs,
                maxBatchSize
            )
        ) {
            acc.appendControlMessages((offset, epoch, compression, buf) -> {
                long now = 1234;
                try (MemoryRecordsBuilder builder = controlRecordsBuilder(
                        offset,
                        epoch,
                        compression,
                        now,
                        buf
                    )
                ) {
                    builder.appendSnapshotHeaderMessage(
                        now,
                        new SnapshotHeaderRecord()
                            .setVersion(ControlRecordUtils.SNAPSHOT_HEADER_CURRENT_VERSION)
                            .setLastContainedLogTimestamp(now)
                    );

                    builder.appendKRaftVersionMessage(
                        now,
                        new KRaftVersionRecord()
                            .setVersion(ControlRecordUtils.KRAFT_VERSION_CURRENT_VERSION)
                            .setKRaftVersion((short) 0)
                    );

                    return builder.build();
                }
            });

            List<BatchAccumulator.CompletedBatch<String>> batches = acc.drain();
            assertEquals(1, batches.size());

            BatchAccumulator.CompletedBatch<String> batch = batches.get(0);
            assertEquals(baseOffset, batch.baseOffset);
            assertEquals(2, batch.numRecords);
            assertEquals(buffer.duplicate().flip(), batch.data.buffer());

            batch.release();
        }
    }

    @Test
    public void testInvalidControlRecordOffset() {
        int leaderEpoch = 17;
        long baseOffset = 157;
        int lingerMs = 50;
        int maxBatchSize = 512;

        ByteBuffer buffer = ByteBuffer.allocate(maxBatchSize);
        Mockito.when(memoryPool.tryAllocate(maxBatchSize))
                .thenReturn(buffer);

        BatchAccumulator.MemoryRecordsCreator creator = (offset, epoch, compression, buf) -> {
            long now = 1234;
            try (MemoryRecordsBuilder builder = controlRecordsBuilder(
                    offset + 1,
                    epoch,
                    compression,
                    now,
                    buf
                )
            ) {
                builder.appendSnapshotHeaderMessage(
                    now,
                    new SnapshotHeaderRecord()
                        .setVersion(ControlRecordUtils.SNAPSHOT_HEADER_CURRENT_VERSION)
                        .setLastContainedLogTimestamp(now)
                );

                return builder.build();
            }
        };

        try (BatchAccumulator<String> acc = buildAccumulator(
                leaderEpoch,
                baseOffset,
                lingerMs,
                maxBatchSize
            )
        ) {
            assertThrows(IllegalArgumentException.class, () -> acc.appendControlMessages(creator));
        }
    }

    @Test
    public void testInvalidControlRecordEpoch() {
        int leaderEpoch = 17;
        long baseOffset = 157;
        int lingerMs = 50;
        int maxBatchSize = 512;

        ByteBuffer buffer = ByteBuffer.allocate(maxBatchSize);
        Mockito.when(memoryPool.tryAllocate(maxBatchSize))
                .thenReturn(buffer);

        BatchAccumulator.MemoryRecordsCreator creator = (offset, epoch, compression, buf) -> {
            long now = 1234;
            try (MemoryRecordsBuilder builder = controlRecordsBuilder(
                    offset,
                    epoch + 1,
                    compression,
                    now,
                    buf
                )
            ) {
                builder.appendSnapshotHeaderMessage(
                    now,
                    new SnapshotHeaderRecord()
                        .setVersion(ControlRecordUtils.SNAPSHOT_HEADER_CURRENT_VERSION)
                        .setLastContainedLogTimestamp(now)
                );

                return builder.build();
            }
        };

        try (BatchAccumulator<String> acc = buildAccumulator(
                leaderEpoch,
                baseOffset,
                lingerMs,
                maxBatchSize
            )
        ) {
            assertThrows(IllegalArgumentException.class, () -> acc.appendControlMessages(creator));
        }
    }

    @Test
    public void testEmptyControlBatch() {
        int leaderEpoch = 17;
        long baseOffset = 157;
        int lingerMs = 50;
        int maxBatchSize = 512;

        ByteBuffer buffer = ByteBuffer.allocate(maxBatchSize);
        Mockito.when(memoryPool.tryAllocate(maxBatchSize))
                .thenReturn(buffer);

        BatchAccumulator.MemoryRecordsCreator creator = (offset, epoch, compression, buf) -> {
            long now = 1234;
            try (MemoryRecordsBuilder builder = controlRecordsBuilder(
                    offset,
                    epoch,
                    compression,
                    now,
                    buf
                )
            ) {
                // Create a control batch without any records
                return builder.build();
            }
        };

        try (BatchAccumulator<String> acc = buildAccumulator(
                leaderEpoch,
                baseOffset,
                lingerMs,
                maxBatchSize
            )
        ) {
            assertThrows(IllegalArgumentException.class, () -> acc.appendControlMessages(creator));
        }
    }

    @Test
    void testDelayedDrain() {
        int leaderEpoch = 17;
        long baseOffset = 157;
        int lingerMs = 50;
        int maxBatchSize = 256;

        Mockito.when(memoryPool.tryAllocate(maxBatchSize))
            .thenReturn(ByteBuffer.allocate(maxBatchSize));

        BatchAccumulator<String> acc = buildAccumulator(
            leaderEpoch,
            baseOffset,
            lingerMs,
            maxBatchSize
        );

        List<String> records = asList("a", "b", "c", "d", "e", "f", "g");
        assertEquals(baseOffset, acc.append(leaderEpoch, records.subList(0, 1), true));
        acc.forceDrain();
        assertEquals(baseOffset + 2, acc.append(leaderEpoch, records.subList(1, 3), true));
        acc.forceDrain();
        assertEquals(baseOffset + 5, acc.append(leaderEpoch, records.subList(3, 6), true));
        acc.forceDrain();

        assertFalse(acc.needsDrain(time.milliseconds()));
        assertNotEquals(0, acc.timeUntilDrain(time.milliseconds()));
        assertEquals(0, acc.drain().size());

        acc.allowDrain();
        assertTrue(acc.needsDrain(time.milliseconds()));
        assertEquals(0, acc.timeUntilDrain(time.milliseconds()));
        assertEquals(3, acc.drain().size());
    }

    @Test
    void testDelayedDrainAreReleased() {
        int leaderEpoch = 17;
        long baseOffset = 157;
        int lingerMs = 50;
        int maxBatchSize = 256;

        ByteBuffer buffer = ByteBuffer.allocate(maxBatchSize);
        Mockito.when(memoryPool.tryAllocate(maxBatchSize))
            .thenReturn(buffer);

        BatchAccumulator<String> acc = buildAccumulator(
            leaderEpoch,
            baseOffset,
            lingerMs,
            maxBatchSize
        );

        List<String> records = asList("a", "b", "c", "d", "e", "f", "g");
        assertEquals(baseOffset, acc.append(leaderEpoch, records.subList(0, 1), true));
        acc.forceDrain();

        acc.close();
        Mockito.verify(memoryPool).release(buffer);
    }

    @Test
    public void testMaxNumberOfBatches() {
        int leaderEpoch = 17;
        long baseOffset = 157;
        int lingerMs = 50;
        int maxBatchSize = 256;
        int maxNumberOfBatches = 2;

        Mockito.when(memoryPool.tryAllocate(maxBatchSize))
            .thenReturn(ByteBuffer.allocate(maxBatchSize));

        StringSerde serde = Mockito.spy(new StringSerde());
        BatchAccumulator<String> acc = new BatchAccumulator<>(
            leaderEpoch,
            baseOffset,
            lingerMs,
            maxBatchSize,
            maxNumberOfBatches,
            memoryPool,
            time,
            Compression.NONE,
            serde
        );

        List<String> records = asList("a", "b", "c", "d", "e", "f", "g");
        assertEquals(baseOffset, acc.append(leaderEpoch, records.subList(0, 1), true));
        acc.forceDrain();
        assertEquals(baseOffset + 2, acc.append(leaderEpoch, records.subList(1, 3), true));
        acc.forceDrain();

        assertThrows(IllegalStateException.class, () -> acc.append(leaderEpoch, records.subList(3, 6), true));
    }

    private static MemoryRecordsBuilder controlRecordsBuilder(
        long baseOffset,
        int epoch,
        Compression compression,
        long now,
        ByteBuffer buffer
    ) {
        return new MemoryRecordsBuilder(
            buffer,
            RecordBatch.CURRENT_MAGIC_VALUE,
            compression,
            TimestampType.CREATE_TIME,
            baseOffset,
            now,
            RecordBatch.NO_PRODUCER_ID,
            RecordBatch.NO_PRODUCER_EPOCH,
            RecordBatch.NO_SEQUENCE,
            false, // isTransactional
            true,  // isControlBatch
            epoch,
            buffer.capacity()
        );
    }
}
