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

Code before:
        // 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);

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

    private void processExactlyOneRecord(StreamsBuilder streamsBuilder, Properties props, String key, String value) throws Exception {
        kafkaStreams = new KafkaStreams(streamsBuilder.build(), props);
        kafkaStreams.start();

        final long currentTimeNew = CLUSTER.time.milliseconds();

        processKeyValueAndVerify(
            key,
            value,
            currentTimeNew,
            List.of(
                new KeyValueTimestamp<>(key, value, currentTimeNew)
            )
        );

        kafkaStreams.close();
        kafkaStreams = null;
    }



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(
            () -> kafkaStreams == null || kafkaStreams.state().isTerminal(),
            10000,
            "KafkaStreams instance did not reach a terminal state within timeout"
        );
```

Code after:
        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());
        processExactlyOneRecord(streamsBuilder, props, "2", "B");

        // 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, "3", "C");
    }

    private void processExactlyOneRecord(
        final StreamsBuilder streamsBuilder,
        final Properties props,
        final String key,
        final String value)
        throws Exception {
        kafkaStreams = new KafkaStreams(streamsBuilder.build(), props);
        kafkaStreams.start();

        final long currentTimeNew = CLUSTER.time.milliseconds();

        processKeyValueAndVerify(
            key,