diff --git a/streams/src/test/java/org/apache/kafka/streams/integration/StreamsUpgradeTestIntegrationTest.java b/streams/src/test/java/org/apache/kafka/streams/integration/StreamsUpgradeTestIntegrationTest.java
index 43c7aa1922e0d..48b459e2a0822 100644
--- a/streams/src/test/java/org/apache/kafka/streams/integration/StreamsUpgradeTestIntegrationTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/integration/StreamsUpgradeTestIntegrationTest.java
@@ -36,6 +36,7 @@
 import static org.apache.kafka.common.utils.Utils.mkMap;
 import static org.apache.kafka.common.utils.Utils.mkProperties;
 import static org.apache.kafka.streams.processor.internals.assignment.StreamsAssignmentProtocolVersions.LATEST_SUPPORTED_VERSION;
+import static org.apache.kafka.test.TestUtils.retryOnExceptionWithTimeout;
 import static org.hamcrest.CoreMatchers.is;
 import static org.hamcrest.MatcherAssert.assertThat;
 
@@ -87,9 +88,9 @@ public void testVersionProbingUpgrade() throws InterruptedException {
         final AtomicInteger usedVersion6 = new AtomicInteger();
         final KafkaStreams kafkaStreams6 = buildFutureStreams(usedVersion6);
         startSync(kafkaStreams6);
-        assertThat(usedVersion6.get(), is(LATEST_SUPPORTED_VERSION + 1));
-        assertThat(usedVersion5.get(), is(LATEST_SUPPORTED_VERSION + 1));
-        assertThat(usedVersion4.get(), is(LATEST_SUPPORTED_VERSION + 1));
+        retryOnExceptionWithTimeout(() -> assertThat(usedVersion6.get(), is(LATEST_SUPPORTED_VERSION + 1)));
+        retryOnExceptionWithTimeout(() -> assertThat(usedVersion5.get(), is(LATEST_SUPPORTED_VERSION + 1)));
+        retryOnExceptionWithTimeout(() -> assertThat(usedVersion4.get(), is(LATEST_SUPPORTED_VERSION + 1)));
 
         kafkaStreams4.close(Duration.ZERO);
         kafkaStreams5.close(Duration.ZERO);
@@ -99,14 +100,14 @@ public void testVersionProbingUpgrade() throws InterruptedException {
         kafkaStreams6.close();
     }
 
-    public static KafkaStreams buildFutureStreams(final AtomicInteger usedVersion4) {
+    private static KafkaStreams buildFutureStreams(final AtomicInteger usedVersion4) {
         final Properties properties = new Properties();
         properties.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, CLUSTER.bootstrapServers());
         properties.put("test.future.metadata", usedVersion4);
         return StreamsUpgradeTest.buildStreams(properties);
     }
 
-    public static void startSync(final KafkaStreams... kafkaStreams) throws InterruptedException {
+    private static void startSync(final KafkaStreams... kafkaStreams) throws InterruptedException {
         final CountDownLatch latch = new CountDownLatch(kafkaStreams.length);
         for (final KafkaStreams streams : kafkaStreams) {
             streams.setStateListener((newState, oldState) -> {