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

Code before:
                new Serdes.StringSerde()
            ),
            "readOnlySource",
            new IntegerDeserializer(),
            new StringDeserializer(),
            inputStream,
            "readOnlyProcessor",
            () -> new ReadOnlyStoreProcessor(numReceived, offsetLimitDelta, shutdownLatch)
        );

        kafkaStreams = new KafkaStreams(topology, props);
        kafkaStreams.setStateListener((newState, oldState) -> {
            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({

Review comment (suggestion):
Probably best to use `close(Duration timeout)` overload, otherwise the test can get stuck

Code after:
                new Serdes.StringSerde()
            ),
            "readOnlySource",
            new IntegerDeserializer(),
            new StringDeserializer(),
            inputStream,
            "readOnlyProcessor",
            () -> new ReadOnlyStoreProcessor(numReceived, offsetLimitDelta, shutdownLatch)
        );

        kafkaStreams = new KafkaStreams(topology, props);

        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());
        }