diff --git a/streams/src/test/java/org/apache/kafka/streams/integration/QueryableStateIntegrationTest.java b/streams/src/test/java/org/apache/kafka/streams/integration/QueryableStateIntegrationTest.java
index 6c17d2df217..67eb4a73d0a 100644
--- a/streams/src/test/java/org/apache/kafka/streams/integration/QueryableStateIntegrationTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/integration/QueryableStateIntegrationTest.java
@@ -573,8 +573,8 @@ public class QueryableStateIntegrationTest {
         final Set<KeyValue<String, Long>> countState = new TreeSet<>(stringLongComparator);
 
         final long timeout = System.currentTimeMillis() + 30000;
-        while (windowState.size() < 5 &&
-            countState.size() < 5 &&
+        while ((windowState.size() < keys.length ||
+            countState.size() < keys.length) &&
             System.currentTimeMillis() < timeout) {
             Thread.sleep(10);
             for (final String key : keys) {