[GitHub] [kafka] showuon commented on a change in pull request #9888: KAFKA-12194: use stateListener to catch each state change

[GitBox <gi...@apache.org>]

showuon commented on a change in pull request #9888:
URL: https://github.com/apache/kafka/pull/9888#discussion_r557446077



##########
File path: streams/src/test/java/org/apache/kafka/streams/integration/AdjustStreamThreadCountTest.java
##########
@@ -88,69 +89,96 @@ public void setup() {
                 mkEntry(StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, Serdes.StringSerde.class)
             )
         );
+
+        createAndRunStream();
     }
 
     @After
     public void teardown() throws IOException {
+        if (kafkaStreams != null) {
+            kafkaStreams.close();
+        }
         purgeLocalStreamsState(properties);
     }
 
+    private void addStreamStateChangeListener(final KafkaStreams kafkaStreams) {
+        // we store each new state in state transition so that we won't miss any state change
+        kafkaStreams.setStateListener(
+            (newState, oldState) -> stateToTransitions.add(newState)
+        );
+    }
+
+    private void waitForStateTransition(final KafkaStreams.State expected) throws InterruptedException {
+        waitForCondition(
+            () -> !stateToTransitions.isEmpty() && stateToTransitions.contains(expected),
+            DEFAULT_DURATION.toMillis(),
+            () -> String.format("Client did not change to the %s state in time. Observed new state transitions: %s",
+                expected, stateToTransitions)
+        );
+    }
+
+    private void createAndRunStream() throws InterruptedException {
+        kafkaStreams = new KafkaStreams(builder.build(), properties);
+        addStreamStateChangeListener(kafkaStreams);
+        kafkaStreams.start();
+        waitForStateTransition(KafkaStreams.State.RUNNING);
+    }
+
     @Test
     public void shouldAddStreamThread() throws Exception {
-        try (final KafkaStreams kafkaStreams = new KafkaStreams(builder.build(), properties)) {
-            StreamsTestUtils.startKafkaStreamsAndWaitForRunningState(kafkaStreams);

Review comment:
       Yes, you're right. I just thought there are some duplicate codes and want to clean them up. But you're right, maybe not better. Fixed.




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org