Repository: apache/kafka
File: coordinator-common/src/test/java/org/apache/kafka/coordinator/common/runtime/CoordinatorRuntimeTest.java (Java)

Code before:
        // Write #2.
        CompletableFuture<String> write2 = runtime.scheduleWriteOperation("write#2", TP, Duration.ofMillis(20),
            state -> new CoordinatorResult<>(List.of("record2"), "response2")
        );

        // The batch should have been flushed.
        assertEquals(1, writer.entries(TP).size());

        // Because flushing the batch cancels the linger task, there should now be two write timeout tasks.
        assertEquals(2, schedulerTimer.size());

        // Verify batch contains both two records
        MemoryRecords batch = writer.entries(TP).get(0);
        RecordBatch recordBatch = batch.firstBatch();
        assertEquals(2, recordBatch.countOrNull());

        // Commit and verify that writes are completed.
        writer.commit(TP);
        assertTrue(write1.isDone());
        assertTrue(write2.isDone());
        // Now that all scheduled tasks have been cancelled, the scheduler queue should be empty.
        assertEquals(0, schedulerTimer.size());
    }

    @Test
    public void testScheduleUnloadingWithPendingBatchWhenPartitionWriterConfigThrows() {
        MockTimer timer = new MockTimer();
        MockPartitionWriter writer = mock(MockPartitionWriter.class);
        MockCoordinatorShardBuilderSupplier supplier = mock(MockCoordinatorShardBuilderSupplier.class);
        MockCoordinatorShardBuilder builder = mock(MockCoordinatorShardBuilder.class);
        MockCoordinatorShard coordinator = mock(MockCoordinatorShard.class);

        CoordinatorRuntime<MockCoordinatorShard, String> runtime =
            new CoordinatorRuntime.Builder<MockCoordinatorShard, String>()
                .withTime(timer.time())
                .withTimer(timer)
                .withDefaultWriteTimeOut(Duration.ofMillis(20))
                .withLoader(new MockCoordinatorLoader())
                .withEventProcessor(new DirectEventProcessor())
                .withPartitionWriter(writer)
                .withCoordinatorShardBuilderSupplier(supplier)
                .withCoordinatorRuntimeMetrics(mock(CoordinatorRuntimeMetrics.class))
                .withCoordinatorMetrics(mock(CoordinatorMetrics.class))
                .withSerializer(new StringSerializer())
                .withAppendLingerMs(OptionalInt.of(10))
                .withExecutorService(mock(ExecutorService.class))
                .build();

        when(builder.withSnapshotRegistry(any())).thenReturn(builder);
        when(builder.withLogContext(any())).thenReturn(builder);
        when(builder.withTime(any())).thenReturn(builder);

Review comment (refactor):
Could we move this together with the rest of the `testScheduleUnloading*` tests?

Code after:

        // Verify the state.
        CoordinatorRuntime<MockCoordinatorShard, String>.CoordinatorContext ctx = runtime.contextOrThrow(TP);
        assertEquals(1L, ctx.coordinator.lastWrittenOffset());
        assertEquals(0L, ctx.coordinator.lastCommittedOffset());
        assertEquals(List.of(
            new MockCoordinatorShard.RecordAndMetadata(0, "record1")
        ), ctx.coordinator.coordinator().fullRecords());

        // Execute the pending write.
        processor.poll();

        // The processor must be empty now.
        assertEquals(0, processor.size());

        // Verify the state.
        assertEquals(2L, ctx.coordinator.lastWrittenOffset());
        assertEquals(0L, ctx.coordinator.lastCommittedOffset());
        assertEquals(List.of(
            new MockCoordinatorShard.RecordAndMetadata(0, "record1"),
            new MockCoordinatorShard.RecordAndMetadata(1, "record2")
        ), ctx.coordinator.coordinator().fullRecords());

        // Commit.
        writer.commit(TP);
        processor.poll();
        assertTrue(write1.isDone());
    }

    @Test
    public void testLingerTimeComparisonInMaybeFlushCurrentBatch() throws Exception {
        // Provides the runtime clock; we will advance it.
        MockTimer clockTimer = new MockTimer();
        // Used for scheduling timer tasks; we won't advance it to avoid a timer-triggered batch flush.
        MockTimer schedulerTimer = new MockTimer();

        MockPartitionWriter writer = new MockPartitionWriter();

        CoordinatorRuntime<MockCoordinatorShard, String> runtime =
            new CoordinatorRuntime.Builder<MockCoordinatorShard, String>()
                .withTime(clockTimer.time())
                .withTimer(schedulerTimer)
                .withDefaultWriteTimeOut(Duration.ofMillis(20))
                .withLoader(new MockCoordinatorLoader())
                .withEventProcessor(new DirectEventProcessor())
                .withPartitionWriter(writer)
                .withCoordinatorShardBuilderSupplier(new MockCoordinatorShardBuilderSupplier())
                .withCoordinatorRuntimeMetrics(mock(CoordinatorRuntimeMetrics.class))
                .withCoordinatorMetrics(mock(CoordinatorMetrics.class))
                .withSerializer(new StringSerializer())
                .withAppendLingerMs(OptionalInt.of(10))