Source:
public void testAwaitFlushCompleteShouldWaitForSplittedBatch() throws Exception { int expectedSplittedBatches = 2; int batchSize = 1025; boolean isSplittedBatch = true; boolean isNonSplittedBatch = false; TopicPartition topicPartition = new TopicPartition("test", 0); MemoryRecordsBuilder memoryRecordsBuilder = MemoryRecords.builder(ByteBuffer.allocate(128), CompressionType.NONE, TimestampType.CREATE_TIME, 128); ProduceRequestResult splittedBatchResult = mock(ProduceRequestResult.class); ProduceRequestResult originalBatchResult = mock(ProduceRequestResult.class); ProducerBatch splittedBatch = new ProducerBatch(topicPartition, memoryRecordsBuilder, 0, isSplittedBatch, splittedBatchResult, new ArrayList<>()); List<ProducerBatch> childrenProducerBatch = new ArrayList<>(); childrenProducerBatch.add(splittedBatch); ProducerBatch originalBatch = new ProducerBatch(topicPartition, memoryRecordsBuilder, 0, isNonSplittedBatch, originalBatchResult, childrenProducerBatch); IncompleteBatches incompleteBatches = new IncompleteBatches(); incompleteBatches.add(splittedBatch); incompleteBatches.add(originalBatch); RecordAccumulator accumulator = createTestRecordAccumulator(3200, batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10L * batchSize, CompressionType.NONE, 10, incompleteBatches); accumulator.awaitFlushCompletion(); verify(originalBatchResult).await(); verify(splittedBatchResult, VerificationModeFactory.times(expectedSplittedBatches)).await(); }

Target:
public void testAwaitFlushCompleteShouldWaitForSplittedBatch() throws Exception { int expectedSplittedBatches = 2; int batchSize = 1025; TopicPartition topicPartition = new TopicPartition("test", 0); MemoryRecordsBuilder memoryRecordsBuilder = MemoryRecords.builder(ByteBuffer.allocate(128), CompressionType.NONE, TimestampType.CREATE_TIME, 128); ProduceRequestResult originalBatchResult = mock(ProduceRequestResult.class); ProduceRequestResult firstSplitBatchResult = mock(ProduceRequestResult.class); ProduceRequestResult secondSplitBatchResult = mock(ProduceRequestResult.class); ProducerBatch secondSplitBatch = new ProducerBatch(topicPartition, memoryRecordsBuilder, 0, true, secondSplitBatchResult, new ArrayList<>()); List<ProducerBatch> grandChildrenProducerBatch = new ArrayList<>(); grandChildrenProducerBatch.add(secondSplitBatch); ProducerBatch firstSplitBatch = new ProducerBatch(topicPartition, memoryRecordsBuilder, 0, true, firstSplitBatchResult, grandChildrenProducerBatch); List<ProducerBatch> childrenProducerBatch = new ArrayList<>(); childrenProducerBatch.add(firstSplitBatch); ProducerBatch originalBatch = new ProducerBatch(topicPartition, memoryRecordsBuilder, 0, false, originalBatchResult, childrenProducerBatch); IncompleteBatches incompleteBatches = new IncompleteBatches(); incompleteBatches.add(secondSplitBatch); incompleteBatches.add(firstSplitBatch); incompleteBatches.add(originalBatch); RecordAccumulator accumulator = createTestRecordAccumulator(3200, batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10L * batchSize, CompressionType.NONE, 10, incompleteBatches); accumulator.awaitFlushCompletion(); verify(originalBatchResult).await(); verify(firstSplitBatchResult, VerificationModeFactory.times(expectedSplittedBatches)).await(); verify(secondSplitBatchResult, VerificationModeFactory.times(expectedSplittedBatches)).await(); }