test: add wait_for_connect to BitcoinTestFramework.connect_nodes

1 files changed, 11 insertions, 1 deletions

diff --git a/test/functional/test_framework/test_framework.py b/test/functional/test_framework/test_framework.py
index ab7fed335c..a34c34713e 100755
--- a/test/functional/test_framework/test_framework.py
+++ b/test/functional/test_framework/test_framework.py
@@ -586,7 +586,14 @@ class BitcoinTestFramework(metaclass=BitcoinTestMetaClass):
     def wait_for_node_exit(self, i, timeout):
         self.nodes[i].process.wait(timeout)
 
-    def connect_nodes(self, a, b, *, peer_advertises_v2=None):
+    def connect_nodes(self, a, b, *, peer_advertises_v2=None, wait_for_connect: bool = True):
+        """
+        Kwargs:
+            wait_for_connect: if True, block until the nodes are verified as connected. You might
+                want to disable this when using -stopatheight with one of the connected nodes,
+                since there will be a race between the actual connection and performing
+                the assertions before one node shuts down.
+        """
         from_connection = self.nodes[a]
         to_connection = self.nodes[b]
         from_num_peers = 1 + len(from_connection.getpeerinfo())
@@ -603,6 +610,9 @@ class BitcoinTestFramework(metaclass=BitcoinTestMetaClass):
             # compatibility with older clients
             from_connection.addnode(ip_port, "onetry")
 
+        if not wait_for_connect:
+            return
+
         # poll until version handshake complete to avoid race conditions
         # with transaction relaying
         # See comments in net_processing: