test: add wait_for_connect to BitcoinTestFramework.connect_nodes

jamesob · jamesob · commit 7005a01c1900 · 2023-10-04T11:05:03.000-04:00
diff --git a/test/functional/test_framework/test_framework.py b/test/functional/test_framework/test_framework.py
@@ -586,7 +586,14 @@ def restart_node(self, i, extra_args=None):
     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 @@ def connect_nodes(self, a, b, *, peer_advertises_v2=None):
             # 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:
