Repository: apache/kafka
File: streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/RebalanceProtocolMigrationIntegrationTest.java (Java)

Code before:
        streamsConfiguration.put(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, 100);
        return streamsConfiguration;
    }

    @AfterEach
    public void shutdown() {
        if (kafkaStreams != null) {
            kafkaStreams.close(Duration.ofSeconds(30L));
            kafkaStreams.cleanUp();
        }
    }


    @Test
    public void shouldMigrateToAndFromStreamsRebalanceProtocol() throws Exception {
        final StreamsBuilder streamsBuilder = new StreamsBuilder();
        final KStream<String, String> input = streamsBuilder.stream(
            inputTopic, Consumed.with(Serdes.String(), Serdes.String()));
        input.to(outputTopic, Produced.with(Serdes.String(), Serdes.String()));

        final Properties props = props();
        props.put(StreamsConfig.GROUP_PROTOCOL_CONFIG, GroupProtocol.CLASSIC.name());
        processExactlyOneRecord(streamsBuilder, props, "1", "A");

        // Wait for session to time out
        Thread.sleep(600);

        props.put(StreamsConfig.GROUP_PROTOCOL_CONFIG, GroupProtocol.STREAMS.name());
        processExactlyOneRecord(streamsBuilder, props, "2", "B");

        props.put(StreamsConfig.GROUP_PROTOCOL_CONFIG, GroupProtocol.CLASSIC.name());
        processExactlyOneRecord(streamsBuilder, props, "3", "C");
    }

    @Test
    public void shouldMigrateFromAndToStreamsRebalanceProtocol() throws Exception {
        final StreamsBuilder streamsBuilder = new StreamsBuilder();
        final KStream<String, String> input = streamsBuilder.stream(
            inputTopic, Consumed.with(Serdes.String(), Serdes.String()));
        input.to(outputTopic, Produced.with(Serdes.String(), Serdes.String()));

        final Properties props = props();
        props.put(StreamsConfig.GROUP_PROTOCOL_CONFIG, GroupProtocol.STREAMS.name());
        processExactlyOneRecord(streamsBuilder, props, "1", "A");

        props.put(StreamsConfig.GROUP_PROTOCOL_CONFIG, GroupProtocol.CLASSIC.name());
        processExactlyOneRecord(streamsBuilder, props, "2", "B");

        // Wait for session to time out
        Thread.sleep(600);


Review comment (suggestion):
Using Thread.sleep() for timing coordination in tests is unreliable and can cause flaky tests. Consider using TestUtils.waitForCondition() or similar mechanisms to wait for the expected state.
```suggestion
        TestUtils.waitForCondition(() -> true, 600, "Wait for session to time out");
```

Code after:
        streamsConfiguration.put(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG, 100L);
        streamsConfiguration.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
        streamsConfiguration.put(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG, 500);
        streamsConfiguration.put(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, 100);
        return streamsConfiguration;
    }

    @AfterEach
    public void shutdown() {
        if (kafkaStreams != null) {
            kafkaStreams.close(Duration.ofSeconds(30L));
            kafkaStreams.cleanUp();
        }
    }


    @Test
    public void shouldMigrateToAndFromStreamsRebalanceProtocol() throws Exception {
        final StreamsBuilder streamsBuilder = new StreamsBuilder();
        final KStream<String, String> input = streamsBuilder.stream(
            inputTopic, Consumed.with(Serdes.String(), Serdes.String()));
        input.to(outputTopic, Produced.with(Serdes.String(), Serdes.String()));

        final Properties props = props();
        props.put(StreamsConfig.GROUP_PROTOCOL_CONFIG, GroupProtocol.CLASSIC.name());
        processExactlyOneRecord(streamsBuilder, props, "1", "A");

        // Wait for session to time out
        try (final Admin adminClient = Admin.create(Map.of(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, CLUSTER.bootstrapServers()))) {
            waitForEmptyConsumerGroup(adminClient, props.getProperty(StreamsConfig.APPLICATION_ID_CONFIG), 1000);
        }

        props.put(StreamsConfig.GROUP_PROTOCOL_CONFIG, GroupProtocol.STREAMS.name());
        processExactlyOneRecord(streamsBuilder, props, "2", "B");

        props.put(StreamsConfig.GROUP_PROTOCOL_CONFIG, GroupProtocol.CLASSIC.name());
        processExactlyOneRecord(streamsBuilder, props, "3", "C");
    }

    @Test
    public void shouldMigrateFromAndToStreamsRebalanceProtocol() throws Exception {
        final StreamsBuilder streamsBuilder = new StreamsBuilder();
        final KStream<String, String> input = streamsBuilder.stream(
            inputTopic, Consumed.with(Serdes.String(), Serdes.String()));
        input.to(outputTopic, Produced.with(Serdes.String(), Serdes.String()));

        final Properties props = props();
        props.put(StreamsConfig.GROUP_PROTOCOL_CONFIG, GroupProtocol.STREAMS.name());
        processExactlyOneRecord(streamsBuilder, props, "1", "A");

        props.put(StreamsConfig.GROUP_PROTOCOL_CONFIG, GroupProtocol.CLASSIC.name());