From 9b24b74f80431f3c81e71fd2de4f062b3ce0d92f Mon Sep 17 00:00:00 2001
From: Marcin Rusek <marcin.rusek@starburstdata.com>
Date: Wed, 11 Dec 2024 16:39:11 +0100
Subject: [PATCH] Add assertConsistently

---
 .../io/trino/testing/assertions/Assert.java   | 21 +++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/testing/trino-testing-services/src/main/java/io/trino/testing/assertions/Assert.java b/testing/trino-testing-services/src/main/java/io/trino/testing/assertions/Assert.java
index 35205085c25c..42cf8c347ce3 100644
--- a/testing/trino-testing-services/src/main/java/io/trino/testing/assertions/Assert.java
+++ b/testing/trino-testing-services/src/main/java/io/trino/testing/assertions/Assert.java
@@ -58,6 +58,27 @@ public static <E extends Exception> void assertEventually(Duration timeout, Dura
         }
     }
 
+    public static <E extends Exception> void assertConsistently(Duration timeout, Duration retryFrequency, Assert.CheckedRunnable<E> assertion)
+            throws E
+    {
+        long start = System.nanoTime();
+        while (!Thread.currentThread().isInterrupted()) {
+            assertion.run();
+
+            if (Duration.nanosSince(start).compareTo(timeout) > 0) {
+                return;
+            }
+
+            try {
+                Thread.sleep(retryFrequency.toMillis());
+            }
+            catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
+                throw new RuntimeException(e);
+            }
+        }
+    }
+
     public interface CheckedRunnable<E extends Exception>
     {
         void run()