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

Code before:
        // 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);
        when(builder.withTimer(any())).thenReturn(builder);
        when(builder.withCoordinatorMetrics(any())).thenReturn(builder);
        when(builder.withTopicPartition(any())).thenReturn(builder);
        when(builder.withExecutor(any())).thenReturn(builder);
        when(builder.build()).thenReturn(coordinator);
        when(supplier.get()).thenReturn(builder);

        // Configure the partition writer with a normal config initially.
        LogConfig initialLogConfig = new LogConfig(
            Map.of(TopicConfig.MAX_MESSAGE_BYTES_CONFIG, String.valueOf(1024 * 1024) // 1MB
        ));

Review comment (suggestion):
nit: Could we use `DEFAULT_WRITE_TIMEOUT` here unless there's a reason to use a different timeout?
```suggestion
                .withDefaultWriteTimeOut(DEFAULT_WRITE_TIMEOUT)
```

Code after:

        // 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))
                .withExecutorService(mock(ExecutorService.class))
                .build();

        // Schedule the loading.
        runtime.scheduleLoadOperation(TP, 10);

        // Verify the initial state.
        CoordinatorRuntime<MockCoordinatorShard, String>.CoordinatorContext ctx = runtime.contextOrThrow(TP);
        assertEquals(ACTIVE, ctx.state);
        assertNull(ctx.currentBatch);