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

Code before:
        // This is the scenario described in KAFKA-20115.
        when(writer.config(TP)).thenThrow(Errors.NOT_LEADER_OR_FOLLOWER.exception());

        // Schedule the unloading. This should trigger the bug where freeCurrentBatch()
        // tries to call partitionWriter.config(tp).maxMessageSize() and throws an exception.
        // Without the fix, this would prevent the coordinator from unloading properly.
        runtime.scheduleUnloadOperation(TP, OptionalInt.of(ctx.epoch + 1));

        // The unload should complete despite the NOT_LEADER_OR_FOLLOWER exception
        // when trying to access partition writer config during buffer cleanup.
        assertEquals(CLOSED, ctx.state);

        // Verify that onUnloaded is called.
        verify(coordinator, times(1)).onUnloaded();

        // Verify that the listener is deregistered.
        verify(writer, times(1)).deregisterListener(
            eq(TP),
            any(PartitionWriter.Listener.class)
        );

        // Getting the coordinator context fails because it no longer exists.
        assertThrows(NotCoordinatorException.class, () -> runtime.contextOrThrow(TP));
    }


    @Test
    public void testScheduleWriteOp() throws ExecutionException, InterruptedException, TimeoutException {
        MockTimer timer = new MockTimer();
        MockPartitionWriter writer = new MockPartitionWriter();

        CoordinatorRuntime<MockCoordinatorShard, String> runtime =
            new CoordinatorRuntime.Builder<MockCoordinatorShard, String>()
                .withTime(timer.time())
                .withTimer(timer)
                .withDefaultWriteTimeOut(DEFAULT_WRITE_TIMEOUT)
                .withLoader(new MockCoordinatorLoader())
                .withEventProcessor(new DirectEventProcessor())
                .withPartitionWriter(writer)
                .withCoordinatorShardBuilderSupplier(new MockCoordinatorShardBuilderSupplier())
                .withCoordinatorRuntimeMetrics(mock(CoordinatorRuntimeMetrics.class))
                .withCoordinatorMetrics(mock(CoordinatorMetrics.class))
                .withSerializer(new StringSerializer())
                .withExecutorService(mock(ExecutorService.class))
                .build();

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

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

Review comment (suggestion):
nit: we have too many newlines here
```suggestion

```

Code after:
        // This is the scenario described in KAFKA-20115.
        when(writer.config(TP)).thenThrow(Errors.NOT_LEADER_OR_FOLLOWER.exception());

        // Schedule the unloading. This should trigger the bug where freeCurrentBatch()
        // tries to call partitionWriter.config(tp).maxMessageSize() and throws an exception.
        // Without the fix, this would prevent the coordinator from unloading properly.
        runtime.scheduleUnloadOperation(TP, OptionalInt.of(ctx.epoch + 1));

        // The unload should complete despite the NOT_LEADER_OR_FOLLOWER exception
        // when trying to access partition writer config during buffer cleanup.
        assertEquals(CLOSED, ctx.state);

        // Verify that onUnloaded is called.
        verify(coordinator, times(1)).onUnloaded();

        // Verify that the listener is deregistered.
        verify(writer, times(1)).deregisterListener(
            eq(TP),
            any(PartitionWriter.Listener.class)
        );

        // Getting the coordinator context fails because it no longer exists.
        assertThrows(NotCoordinatorException.class, () -> runtime.contextOrThrow(TP));
    }

    @Test
    public void testScheduleWriteOp() throws ExecutionException, InterruptedException, TimeoutException {
        MockTimer timer = new MockTimer();
        MockPartitionWriter writer = new MockPartitionWriter();

        CoordinatorRuntime<MockCoordinatorShard, String> runtime =
            new CoordinatorRuntime.Builder<MockCoordinatorShard, String>()
                .withTime(timer.time())
                .withTimer(timer)
                .withDefaultWriteTimeOut(DEFAULT_WRITE_TIMEOUT)
                .withLoader(new MockCoordinatorLoader())
                .withEventProcessor(new DirectEventProcessor())
                .withPartitionWriter(writer)
                .withCoordinatorShardBuilderSupplier(new MockCoordinatorShardBuilderSupplier())
                .withCoordinatorRuntimeMetrics(mock(CoordinatorRuntimeMetrics.class))
                .withCoordinatorMetrics(mock(CoordinatorMetrics.class))
                .withSerializer(new StringSerializer())
                .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(0L, ctx.coordinator.lastWrittenOffset());