From 0470e2bc9585acb64eb19a64c4099e495890d355 Mon Sep 17 00:00:00 2001
From: "A. Sophie Blee-Goldman" <sophie@confluent.io>
Date: Thu, 9 Apr 2020 15:28:44 -0700
Subject: [PATCH] KAFKA-6145: KIP-441: fix flaky shouldEnforceRebalance test in
 StreamThreadTest (#8452)

Reviewers: Boyang Chen <boyang@confluent.io>, John Roesler <vvcephei@apache.org>
---
 .../kafka/streams/processor/internals/StreamThreadTest.java  | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamThreadTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamThreadTest.java
index dbd43488bfe..d0488f62d5b 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamThreadTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamThreadTest.java
@@ -508,13 +508,16 @@ public class StreamThreadTest {
             10 * 1000,
             "Thread never started.");
 
+        TestUtils.retryOnExceptionWithTimeout(
+            () -> EasyMock.verify(mockConsumer)
+        );
+
         thread.shutdown();
         TestUtils.waitForCondition(
             () -> thread.state() == StreamThread.State.DEAD,
             10 * 1000,
             "Thread never shut down.");
 
-        EasyMock.verify(mockConsumer);
     }
 
     private static class EasyMockConsumerClientSupplier extends MockClientSupplier {