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

Code before:
        }

        // Write #1.
        CompletableFuture<String> write1 = runtime.scheduleWriteOperation("write#1", TP, DEFAULT_WRITE_TIMEOUT,
            state -> new CoordinatorResult<>(records, "response1")
        );

        // Verify that the write has not completed exceptionally.
        // This will catch any exceptions thrown including RecordTooLargeException.
        assertFalse(write1.isCompletedExceptionally());

        // Verify that the next buffer retrieved from the bufferSupplier is the expanded buffer.
        assertTrue(ctx.bufferSupplier.get(1).capacity() > INITIAL_BUFFER_SIZE);
    }

    @Test
    public void testBufferShrinkWhenMaxMessageSizeReducedBelowInitialBufferSize() {
        MockTimer timer = new MockTimer();
        var mockWriter = new InMemoryPartitionWriter(false) {
            private LogConfig config;

            {
                config = new LogConfig(Map.of(
                    TopicConfig.MAX_MESSAGE_BYTES_CONFIG, String.valueOf(1024 * 1024) // 1MB
                ));
            }

            @Override
            public LogConfig config(TopicPartition tp) {
                return config;
            }

            public void updateConfig(LogConfig newConfig) {
                this.config = newConfig;
            }
        };
        StringSerializer serializer = new StringSerializer();

        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(mockWriter)
                .withCoordinatorShardBuilderSupplier(new MockCoordinatorShardBuilderSupplier())
                .withCoordinatorRuntimeMetrics(mock(CoordinatorRuntimeMetrics.class))
                .withCoordinatorMetrics(mock(CoordinatorMetrics.class))
                .withSerializer(serializer)
                .withExecutorService(mock(ExecutorService.class))

Review comment (suggestion):
How about 
```suggestion
            private LogConfig config = new LogConfig(Map.of(
                TopicConfig.MAX_MESSAGE_BYTES_CONFIG, String.valueOf(1024 * 1024) // 1MB
            ));
```

Code after:
        }

        // Write #1.
        CompletableFuture<String> write1 = runtime.scheduleWriteOperation("write#1", TP, DEFAULT_WRITE_TIMEOUT,
            state -> new CoordinatorResult<>(records, "response1")
        );

        // Verify that the write has not completed exceptionally.
        // This will catch any exceptions thrown including RecordTooLargeException.
        assertFalse(write1.isCompletedExceptionally());

        int batchSize = mockWriter.entries(TP).get(0).sizeInBytes();
        int maxBatchSize = mockWriter.config(TP).maxMessageSize();
        assertTrue(INITIAL_BUFFER_SIZE < batchSize && batchSize <= maxBatchSize);

        // Verify that the next buffer retrieved from the bufferSupplier is the expanded buffer.
        assertTrue(ctx.bufferSupplier.get(1).capacity() > INITIAL_BUFFER_SIZE);
    }

    @Test
    public void testBufferShrinkWhenMaxMessageSizeReducedBelowInitialBufferSize() {
        MockTimer timer = new MockTimer();
        var mockWriter = new InMemoryPartitionWriter(false) {
            private LogConfig config = new LogConfig(Map.of(
                TopicConfig.MAX_MESSAGE_BYTES_CONFIG, String.valueOf(1024 * 1024) // 1MB
            ));

            @Override
            public LogConfig config(TopicPartition tp) {
                return config;
            }

            public void updateConfig(LogConfig newConfig) {
                this.config = newConfig;
            }
        };
        StringSerializer serializer = new StringSerializer();

        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(mockWriter)
                .withCoordinatorShardBuilderSupplier(new MockCoordinatorShardBuilderSupplier())
                .withCoordinatorRuntimeMetrics(mock(CoordinatorRuntimeMetrics.class))
                .withCoordinatorMetrics(mock(CoordinatorMetrics.class))
                .withSerializer(serializer)
                .withExecutorService(mock(ExecutorService.class))