GH-2222: Fix Race in Test

Author: garyrussell

spring-kafka/src/test/java/org/springframework/kafka/listener/KafkaMessageListenerContainerTests.java

@@ -2544,17 +2544,25 @@ public void rePausePartitionAfterRebalance() throws Exception {
 		TopicPartition tp1 = new TopicPartition("foo", 1);
 		given(consumer.assignment()).willReturn(Set.of(tp0, tp1));
 		final CountDownLatch pauseLatch1 = new CountDownLatch(1);
-		final CountDownLatch pauseLatch2 = new CountDownLatch(2);
+		final CountDownLatch suspendConsumerThread = new CountDownLatch(1);
 		Set<TopicPartition> pausedParts = ConcurrentHashMap.newKeySet();
+		Thread testThread = Thread.currentThread();
+		AtomicBoolean paused = new AtomicBoolean();
 		willAnswer(i -> {
 			pausedParts.clear();
 			pausedParts.addAll(i.getArgument(0));
-			pauseLatch1.countDown();
-			pauseLatch2.countDown();
+			if (!Thread.currentThread().equals(testThread)) {
+				paused.set(true);
+			}
 			return null;
 		}).given(consumer).pause(any());
 		given(consumer.paused()).willReturn(pausedParts);
 		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);
+			}
 			Thread.sleep(50);
 			return ConsumerRecords.empty();
 		});
@@ -2585,15 +2593,14 @@ public void rePausePartitionAfterRebalance() throws Exception {
 				.contains(tp0, tp1);
 		rebal.get().onPartitionsRevoked(Set.of(tp0, tp1));
 		rebal.get().onPartitionsAssigned(Collections.singleton(tp0));
-		assertThat(pauseLatch2.await(10, TimeUnit.SECONDS)).isTrue();
 		assertThat(pausedParts).hasSize(1)
 				.contains(tp0);
 		assertThat(container).extracting("listenerConsumer")
 				.extracting("pausedPartitions")
 				.asInstanceOf(InstanceOfAssertFactories.collection(TopicPartition.class))
 				.hasSize(1)
-				.containsExactlyInAnyOrder(tp0);
-
+				.contains(tp0);
+		suspendConsumerThread.countDown();
 		container.stop();
 	}