qa: Move wait_until to util

3 files changed, 35 insertions, 32 deletions

diff --git a/test/functional/test_framework/comptool.py b/test/functional/test_framework/comptool.py
index 9f062865a3..bfbc0c3b03 100755
--- a/test/functional/test_framework/comptool.py
+++ b/test/functional/test_framework/comptool.py
@@ -19,7 +19,7 @@ TestNode behaves as follows:
 
 from .mininode import *
 from .blockstore import BlockStore, TxStore
-from .util import p2p_port
+from .util import p2p_port, wait_until
 
 import logging
 
@@ -189,7 +189,7 @@ class TestManager(object):
     def wait_for_disconnections(self):
         def disconnected():
             return all(node.closed for node in self.test_nodes)
-        return wait_until(disconnected, timeout=10)
+        wait_until(disconnected, timeout=10, lock=mininode_lock)
 
     def wait_for_verack(self):
         return all(node.wait_for_verack() for node in self.test_nodes)
@@ -197,7 +197,7 @@ class TestManager(object):
     def wait_for_pings(self, counter):
         def received_pongs():
             return all(node.received_ping_response(counter) for node in self.test_nodes)
-        return wait_until(received_pongs)
+        wait_until(received_pongs, lock=mininode_lock)
 
     # sync_blocks: Wait for all connections to request the blockhash given
     # then send get_headers to find out the tip of each node, and synchronize
@@ -210,8 +210,7 @@ class TestManager(object):
             )
 
         # --> error if not requested
-        if not wait_until(blocks_requested, attempts=20*num_blocks):
-            raise AssertionError("Not all nodes requested block")
+        wait_until(blocks_requested, attempts=20*num_blocks, lock=mininode_lock)
 
         # Send getheaders message
         [ c.cb.send_getheaders() for c in self.connections ]
@@ -231,8 +230,7 @@ class TestManager(object):
             )
 
         # --> error if not requested
-        if not wait_until(transaction_requested, attempts=20*num_events):
-            raise AssertionError("Not all nodes requested transaction")
+        wait_until(transaction_requested, attempts=20*num_events, lock=mininode_lock)
 
         # Get the mempool
         [ c.cb.send_mempool() for c in self.connections ]
diff --git a/test/functional/test_framework/mininode.py b/test/functional/test_framework/mininode.py
index a4d85501e7..d0753276db 100755
--- a/test/functional/test_framework/mininode.py
+++ b/test/functional/test_framework/mininode.py
@@ -35,7 +35,7 @@ import time
 from threading import RLock, Thread
 
 from test_framework.siphash import siphash256
-from test_framework.util import hex_str_to_bytes, bytes_to_hex_str
+from test_framework.util import hex_str_to_bytes, bytes_to_hex_str, wait_until
 
 BIP0031_VERSION = 60000
 MY_VERSION = 70014  # past bip-31 for ping/pong
@@ -1358,23 +1358,6 @@ class msg_reject(object):
         return "msg_reject: %s %d %s [%064x]" \
             % (self.message, self.code, self.reason, self.data)
 
-# Helper function
-def wait_until(predicate, *, attempts=float('inf'), timeout=float('inf')):
-    if attempts == float('inf') and timeout == float('inf'):
-        timeout = 60
-    attempt = 0
-    elapsed = 0
-
-    while attempt < attempts and elapsed < timeout:
-        with mininode_lock:
-            if predicate():
-                return True
-        attempt += 1
-        elapsed += 0.05
-        time.sleep(0.05)
-
-    return False
-
 class msg_feefilter(object):
     command = b"feefilter"
 
@@ -1591,21 +1574,21 @@ class NodeConnCB(object):
 
     def wait_for_disconnect(self, timeout=60):
         test_function = lambda: not self.connected
-        assert wait_until(test_function, timeout=timeout)
+        wait_until(test_function, timeout=timeout, lock=mininode_lock)
 
     # Message receiving helper methods
 
     def wait_for_block(self, blockhash, timeout=60):
         test_function = lambda: self.last_message.get("block") and self.last_message["block"].block.rehash() == blockhash
-        assert wait_until(test_function, timeout=timeout)
+        wait_until(test_function, timeout=timeout, lock=mininode_lock)
 
     def wait_for_getdata(self, timeout=60):
         test_function = lambda: self.last_message.get("getdata")
-        assert wait_until(test_function, timeout=timeout)
+        wait_until(test_function, timeout=timeout, lock=mininode_lock)
 
     def wait_for_getheaders(self, timeout=60):
         test_function = lambda: self.last_message.get("getheaders")
-        assert wait_until(test_function, timeout=timeout)
+        wait_until(test_function, timeout=timeout, lock=mininode_lock)
 
     def wait_for_inv(self, expected_inv, timeout=60):
         """Waits for an INV message and checks that the first inv object in the message was as expected."""
@@ -1614,11 +1597,11 @@ class NodeConnCB(object):
         test_function = lambda: self.last_message.get("inv") and \
                                 self.last_message["inv"].inv[0].type == expected_inv[0].type and \
                                 self.last_message["inv"].inv[0].hash == expected_inv[0].hash
-        assert wait_until(test_function, timeout=timeout)
+        wait_until(test_function, timeout=timeout, lock=mininode_lock)
 
     def wait_for_verack(self, timeout=60):
         test_function = lambda: self.message_count["verack"]
-        assert wait_until(test_function, timeout=timeout)
+        wait_until(test_function, timeout=timeout, lock=mininode_lock)
 
     # Message sending helper functions
 
@@ -1636,7 +1619,7 @@ class NodeConnCB(object):
     def sync_with_ping(self, timeout=60):
         self.send_message(msg_ping(nonce=self.ping_counter))
         test_function = lambda: self.last_message.get("pong") and self.last_message["pong"].nonce == self.ping_counter
-        assert wait_until(test_function, timeout=timeout)
+        wait_until(test_function, timeout=timeout, lock=mininode_lock)
         self.ping_counter += 1
         return True
 
diff --git a/test/functional/test_framework/util.py b/test/functional/test_framework/util.py
index 4098fd8615..a14cda07d0 100644
--- a/test/functional/test_framework/util.py
+++ b/test/functional/test_framework/util.py
@@ -157,6 +157,28 @@ def str_to_b64str(string):
 def satoshi_round(amount):
     return Decimal(amount).quantize(Decimal('0.00000001'), rounding=ROUND_DOWN)
 
+def wait_until(predicate, *, attempts=float('inf'), timeout=float('inf'), lock=None):
+    if attempts == float('inf') and timeout == float('inf'):
+        timeout = 60
+    attempt = 0
+    timeout += time.time()
+
+    while attempt < attempts and time.time() < timeout:
+        if lock:
+            with lock:
+                if predicate():
+                    return
+        else:
+            if predicate():
+                return
+        attempt += 1
+        time.sleep(0.05)
+
+    # Print the cause of the timeout
+    assert_greater_than(attempts, attempt)
+    assert_greater_than(timeout, time.time())
+    raise RuntimeError('Unreachable')
+
 # RPC/P2P connection constants and functions
 ############################################