From d2e213a1d3fdc8d2d205d35c0d0aea1321b1d421 Mon Sep 17 00:00:00 2001
From: Gary Russell <grussell@vmware.com>
Date: Wed, 25 Oct 2023 14:11:15 -0400
Subject: [PATCH] Fix Race in Test

---
 .../kafka/listener/KafkaMessageListenerContainerTests.java     | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/spring-kafka/src/test/java/org/springframework/kafka/listener/KafkaMessageListenerContainerTests.java b/spring-kafka/src/test/java/org/springframework/kafka/listener/KafkaMessageListenerContainerTests.java
index 05dce516..3d706c87 100644
--- a/spring-kafka/src/test/java/org/springframework/kafka/listener/KafkaMessageListenerContainerTests.java
+++ b/spring-kafka/src/test/java/org/springframework/kafka/listener/KafkaMessageListenerContainerTests.java
@@ -2790,12 +2790,14 @@ public class KafkaMessageListenerContainerTests {
 			return null;
 		}).given(consumer).pause(any());
 		given(consumer.paused()).willReturn(pausedParts);
+		CountDownLatch firstPoll = new CountDownLatch(1);
 		given(consumer.poll(any(Duration.class))).willAnswer(i -> {
 			if (paused.get()) {
 				pauseLatch1.countDown();
 				// hold up the consumer thread while we revoke/assign partitions on the test thread
 				suspendConsumerThread.await(10, TimeUnit.SECONDS);
 			}
+			firstPoll.countDown();
 			Thread.sleep(50);
 			return ConsumerRecords.empty();
 		});
@@ -2819,6 +2821,7 @@ public class KafkaMessageListenerContainerTests {
 				new KafkaMessageListenerContainer<>(cf, containerProps);
 		container.start();
 		InOrder inOrder = inOrder(consumer);
+		assertThat(firstPoll.await(10, TimeUnit.SECONDS)).isNotNull();
 		container.pausePartition(tp0);
 		container.pausePartition(tp1);
 		assertThat(pauseLatch1.await(10, TimeUnit.SECONDS)).isTrue();