MINOR: Use `pollUntilTrue` instead of `waitForCondition` (#19911)

We can use `pollUntilTrue` instead of `waitForCondition`, thus do a
little refactor to reduce the duplicate code

Reviewers: TengYao Chi <frankvicky@apache.org>, Lan Ding
 <isDing_L@163.com>, TaiJuWu <tjwu1217@gmail.com>

clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/ClientsTestUtils.java

@@ -109,12 +109,31 @@ public class ClientsTestUtils {
     }
 
     public static void pollUntilTrue(
-            Consumer<byte[], byte[]> consumer,
+        Consumer<byte[], byte[]> consumer,
-            Supplier<Boolean> testCondition,
+        Supplier<Boolean> testCondition,
-            long waitTimeMs, String msg
+        String msg
+    ) throws InterruptedException {
+        pollUntilTrue(consumer, Duration.ofMillis(100), testCondition, 15_000L, msg);
+    }
+
+    public static void pollUntilTrue(
+        Consumer<byte[], byte[]> consumer,
+        Supplier<Boolean> testCondition,
+        long waitTimeMs,
+        String msg
+    ) throws InterruptedException {
+        pollUntilTrue(consumer, Duration.ofMillis(100), testCondition, waitTimeMs, msg);
+    }
+
+    public static void pollUntilTrue(
+        Consumer<byte[], byte[]> consumer,
+        Duration timeout,
+        Supplier<Boolean> testCondition,
+        long waitTimeMs, 
+        String msg
     ) throws InterruptedException {
         TestUtils.waitForCondition(() -> {
-            consumer.poll(Duration.ofMillis(100));
+            consumer.poll(timeout);
             return testCondition.get();
         }, waitTimeMs, msg);
     }

clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/PlaintextConsumerCallbackTest.java

@@ -16,14 +16,13 @@
  */
 package org.apache.kafka.clients.consumer;
 
+import org.apache.kafka.clients.ClientsTestUtils;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.test.ClusterInstance;
-import org.apache.kafka.common.test.TestUtils;
 import org.apache.kafka.common.test.api.ClusterTest;
 import org.apache.kafka.common.test.api.ClusterTestDefaults;
 import org.apache.kafka.common.test.api.Type;
 
-import java.time.Duration;
 import java.util.Collection;
 import java.util.List;
 import java.util.Locale;
@@ -237,11 +236,9 @@ public class PlaintextConsumerCallbackTest {
                 // noop
             }
         });
-        TestUtils.waitForCondition(
+        ClientsTestUtils.pollUntilTrue(
-            () -> {
+            consumer, 
-                consumer.poll(Duration.ofMillis(100));
+            partitionsAssigned::get, 
-                return partitionsAssigned.get();
-            },
             "Timed out before expected rebalance completed"
         );
     }
@@ -272,11 +269,9 @@ public class PlaintextConsumerCallbackTest {
                     }
                 }
             });
-            TestUtils.waitForCondition(
+            ClientsTestUtils.pollUntilTrue(
-                () -> {
+                consumer,
-                    consumer.poll(Duration.ofMillis(100));
+                partitionsAssigned::get,
-                    return partitionsAssigned.get();
-                },
                 "Timed out before expected rebalance completed"
             );
         }

clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/PlaintextConsumerCommitTest.java

@@ -16,13 +16,13 @@
  */
 package org.apache.kafka.clients.consumer;
 
+import org.apache.kafka.clients.ClientsTestUtils;
 import org.apache.kafka.clients.producer.Producer;
 import org.apache.kafka.clients.producer.ProducerConfig;
 import org.apache.kafka.clients.producer.ProducerRecord;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.internals.Topic;
 import org.apache.kafka.common.test.ClusterInstance;
-import org.apache.kafka.common.test.TestUtils;
 import org.apache.kafka.common.test.api.ClusterConfigProperty;
 import org.apache.kafka.common.test.api.ClusterTest;
 import org.apache.kafka.common.test.api.ClusterTestDefaults;
@@ -203,10 +203,11 @@ public class PlaintextConsumerCommitTest {
             for (var i = 1; i <= count; i++)
                 consumer.commitAsync(Map.of(tp, new OffsetAndMetadata(i)), callback);
 
-            TestUtils.waitForCondition(() -> {
+            ClientsTestUtils.pollUntilTrue(
-                consumer.poll(Duration.ofMillis(100));
+                consumer,
-                return callback.successCount >= count || callback.lastError.isPresent();
+                () -> callback.successCount >= count || callback.lastError.isPresent(),
-            }, "Failed to observe commit callback before timeout");
+                "Failed to observe commit callback before timeout"
+            );
 
             assertEquals(Optional.empty(), callback.lastError);
             assertEquals(count, callback.successCount);
@@ -533,10 +534,10 @@ public class PlaintextConsumerCommitTest {
         var commitCallback = new RetryCommitCallback(consumer, offsetsOpt);
 
         commitCallback.sendAsyncCommit();
-        TestUtils.waitForCondition(() -> {
+        ClientsTestUtils.pollUntilTrue(
-                consumer.poll(Duration.ofMillis(100));
+            consumer,
-                return commitCallback.isComplete;
+            () -> commitCallback.isComplete,
-            }, "Failed to observe commit callback before timeout"
+            "Failed to observe commit callback before timeout"
         );
 
         assertEquals(Optional.empty(), commitCallback.error);

clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/PlaintextConsumerPollTest.java

@@ -17,6 +17,7 @@
 package org.apache.kafka.clients.consumer;
 
 
+import org.apache.kafka.clients.ClientsTestUtils;
 import org.apache.kafka.clients.ClientsTestUtils.TestConsumerReassignmentListener;
 import org.apache.kafka.clients.producer.Producer;
 import org.apache.kafka.common.MetricName;
@@ -559,10 +560,11 @@ public class PlaintextConsumerPollTest {
 
             // Subscribe to different topic. This will trigger the delayed revocation exceeding rebalance timeout and get fenced
             consumer.subscribe(List.of(otherTopic), listener);
-            TestUtils.waitForCondition(() -> {
+            ClientsTestUtils.pollUntilTrue(
-                consumer.poll(Duration.ofMillis(100));
+                consumer,
-                return rebalanceTimeoutExceeded.get();
+                rebalanceTimeoutExceeded::get,
-            }, "Timeout waiting for delayed callback to complete");
+                "Timeout waiting for delayed callback to complete"
+            );
 
             // Verify consumer recovers after being fenced, being able to continue consuming.
             // (The member should automatically rejoin on the next poll, with the new topic as subscription)