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

Code before:
            if (newState == KafkaStreams.State.RUNNING && oldState == KafkaStreams.State.REBALANCING) {
                startupLatch.countDown();
            }
        });

        final AtomicLong restored = new AtomicLong(0);
        kafkaStreams.setGlobalStateRestoreListener(new TrackingStateRestoreListener(restored));
        kafkaStreams.start();

        assertTrue(startupLatch.await(30, TimeUnit.SECONDS));

        if (useNewProtocol) {
            // For new protocol, we need to stop the streams instance before altering offsets
            kafkaStreams.close();
            setCommittedOffset(inputStream, offsetLimitDelta, useNewProtocol);
            
            // Restart the streams instance with a new startup latch
            final CountDownLatch restartLatch = new CountDownLatch(1);
            kafkaStreams = new KafkaStreams(topology, props);
            kafkaStreams.setStateListener((newState, oldState) -> {
                if (newState == KafkaStreams.State.RUNNING && oldState == KafkaStreams.State.REBALANCING) {
                    restartLatch.countDown();
                }
            });
            kafkaStreams.setGlobalStateRestoreListener(new TrackingStateRestoreListener(restored));
            kafkaStreams.start();
            
            // Wait for the restarted instance to be running
            assertTrue(restartLatch.await(30, TimeUnit.SECONDS));
        }

        assertThat(restored.get(), equalTo((long) numberOfKeys - offsetLimitDelta * 2 - offsetCheckpointed * 2));

        assertTrue(shutdownLatch.await(30, TimeUnit.SECONDS));
        assertThat(numReceived.get(), equalTo(offsetLimitDelta * 2));
    }

    @ParameterizedTest
    @CsvSource({
        "true,true",
        "true,false",
        "false,true",
        "false,false"
    })
    public void shouldRestoreStateFromSourceTopicForGlobalTable(final boolean stateUpdaterEnabled, final boolean useNewProtocol) throws Exception {
        final AtomicInteger numReceived = new AtomicInteger(0);
        final StreamsBuilder builder = new StreamsBuilder();

        final Properties props = props(stateUpdaterEnabled);
        props.put(StreamsConfig.TOPOLOGY_OPTIMIZATION_CONFIG, StreamsConfig.OPTIMIZE);
        if (useNewProtocol) {

Review comment (refactor):
There's a ` IntegrationTestUtils.startApplicationAndWaitUntilRunning(KafkaStreams instance)` method you can use that returns your KS instance once it's in a running state which will simplify your logic here - no need for `CountDownLatch`

Code after:
        final AtomicLong restored = new AtomicLong(0);
        kafkaStreams.setGlobalStateRestoreListener(new TrackingStateRestoreListener(restored));
        startApplicationAndWaitUntilRunning(kafkaStreams);

        if (useNewProtocol) {
            // For new protocol, we need to stop the streams instance before altering offsets
            kafkaStreams.close(Duration.ofSeconds(60));
            setCommittedOffset(inputStream, offsetLimitDelta, useNewProtocol);
            
            // Restart the streams instance with a new startup latch
            
            kafkaStreams = new KafkaStreams(topology, props);
            kafkaStreams.setGlobalStateRestoreListener(new TrackingStateRestoreListener(restored));
            startApplicationAndWaitUntilRunning(kafkaStreams);
        }

        assertThat(restored.get(), equalTo((long) numberOfKeys - offsetLimitDelta * 2 - offsetCheckpointed * 2));

        assertTrue(shutdownLatch.await(30, TimeUnit.SECONDS));
        assertThat(numReceived.get(), equalTo(offsetLimitDelta * 2));
    }

    @ParameterizedTest
    @CsvSource({
        "true,true",
        "true,false",
        "false,true",
        "false,false"
    })
    public void shouldRestoreStateFromSourceTopicForGlobalTable(final boolean stateUpdaterEnabled, final boolean useNewProtocol) throws Exception {
        final AtomicInteger numReceived = new AtomicInteger(0);
        final StreamsBuilder builder = new StreamsBuilder();

        final Properties props = props(stateUpdaterEnabled);
        props.put(StreamsConfig.TOPOLOGY_OPTIMIZATION_CONFIG, StreamsConfig.OPTIMIZE);
        if (useNewProtocol) {
            props.put(StreamsConfig.GROUP_PROTOCOL_CONFIG, GroupProtocol.STREAMS.name());
        }

        // restoring from 1000 to 4000 (committed), and then process from 4000 to 5000 on each of the two partitions
        final int offsetLimitDelta = 1000;
        final int offsetCheckpointed = 1000;
        createStateForRestoration(inputStream, 0);
        if (!useNewProtocol) {
            setCommittedOffset(inputStream, offsetLimitDelta, useNewProtocol);
        }

        final StateDirectory stateDirectory = new StateDirectory(new StreamsConfig(props), new MockTime(), true, false);
        // note here the checkpointed offset is the last processed record's offset, so without control message we should write this offset - 1
        new OffsetCheckpoint(new File(stateDirectory.getOrCreateDirectoryForTask(new TaskId(0, 0)), ".checkpoint"))
            .write(Collections.singletonMap(new TopicPartition(inputStream, 0), (long) offsetCheckpointed - 1));