From 1e7653e08a3778446ff677bb147df68b734a31fd Mon Sep 17 00:00:00 2001
From: jp1ac4 <121959000+jp1ac4@users.noreply.github.com>
Date: Thu, 8 Feb 2024 13:11:25 +0000
Subject: [PATCH] tests: add function to wait while condition holds

This adds a new utils function that is a slight generalisation
of `wait_for` with an additional condition that must always be
met while waiting.

`wait_for` now calls this new function with the condition being
one that is always true.
---
 tests/test_framework/utils.py | 22 +++++++++++++++++++---
 1 file changed, 19 insertions(+), 3 deletions(-)

diff --git a/tests/test_framework/utils.py b/tests/test_framework/utils.py
index 8829fd5e..4aabf759 100644
--- a/tests/test_framework/utils.py
+++ b/tests/test_framework/utils.py
@@ -35,17 +35,33 @@ def wait_for(success, timeout=TIMEOUT, debug_fn=None):
     debug_fn is logged at each call to success, it can be useful for debugging
     when tests fail.
     """
+    wait_for_while_condition_holds(success, lambda: True, timeout, debug_fn)
+
+
+def wait_for_while_condition_holds(success, condition, timeout=TIMEOUT, debug_fn=None):
+    """
+    Run success() either until it returns True, or until the timeout is reached,
+    as long as condition() holds.
+    debug_fn is logged at each call to success, it can be useful for debugging
+    when tests fail.
+    """
     start_time = time.time()
     interval = 0.25
-    while not success() and time.time() < start_time + timeout:
+    while True:
+        if time.time() >= start_time + timeout:
+            raise ValueError("Error waiting for {}", success)
+        if not condition():
+            raise ValueError(
+                "Condition {} not met while waiting for {}", condition, success
+            )
+        if success():
+            return
         if debug_fn is not None:
             logging.info(debug_fn())
         time.sleep(interval)
         interval *= 2
         if interval > 5:
             interval = 5
-    if time.time() > start_time + timeout:
-        raise ValueError("Error waiting for {}", success)
 
 
 def get_txid(hex_tx):