1 files changed, 10 insertions, 11 deletions
diff --git a/test/functional/p2p_tx_download.py b/test/functional/p2p_tx_download.py
index 3ea1c6e5e7..653c7ae43f 100755
--- a/test/functional/p2p_tx_download.py
+++ b/test/functional/p2p_tx_download.py
@@ -16,14 +16,13 @@ from test_framework.messages import (
     msg_inv,
     msg_notfound,
 )
-from test_framework.mininode import (
+from test_framework.p2p import (
     P2PInterface,
-    mininode_lock,
+    p2p_lock,
 )
 from test_framework.test_framework import BitcoinTestFramework
 from test_framework.util import (
     assert_equal,
-    wait_until,
 )
 from test_framework.address import ADDRESS_BCRT1_UNSPENDABLE
 
@@ -73,14 +72,14 @@ class TxDownloadTest(BitcoinTestFramework):
 
         def getdata_found(peer_index):
             p = self.nodes[0].p2ps[peer_index]
-            with mininode_lock:
+            with p2p_lock:
                 return p.last_message.get("getdata") and p.last_message["getdata"].inv[-1].hash == txid
 
         node_0_mocktime = int(time.time())
         while outstanding_peer_index:
             node_0_mocktime += MAX_GETDATA_INBOUND_WAIT
             self.nodes[0].setmocktime(node_0_mocktime)
-            wait_until(lambda: any(getdata_found(i) for i in outstanding_peer_index))
+            self.wait_until(lambda: any(getdata_found(i) for i in outstanding_peer_index))
             for i in outstanding_peer_index:
                 if getdata_found(i):
                     outstanding_peer_index.remove(i)
@@ -134,24 +133,24 @@ class TxDownloadTest(BitcoinTestFramework):
 
         p = self.nodes[0].p2ps[0]
 
-        with mininode_lock:
+        with p2p_lock:
             p.tx_getdata_count = 0
 
         p.send_message(msg_inv([CInv(t=MSG_WTX, h=i) for i in txids]))
-        wait_until(lambda: p.tx_getdata_count >= MAX_GETDATA_IN_FLIGHT, lock=mininode_lock)
-        with mininode_lock:
+        p.wait_until(lambda: p.tx_getdata_count >= MAX_GETDATA_IN_FLIGHT)
+        with p2p_lock:
             assert_equal(p.tx_getdata_count, MAX_GETDATA_IN_FLIGHT)
 
         self.log.info("Now check that if we send a NOTFOUND for a transaction, we'll get one more request")
         p.send_message(msg_notfound(vec=[CInv(t=MSG_WTX, h=txids[0])]))
-        wait_until(lambda: p.tx_getdata_count >= MAX_GETDATA_IN_FLIGHT + 1, timeout=10, lock=mininode_lock)
-        with mininode_lock:
+        p.wait_until(lambda: p.tx_getdata_count >= MAX_GETDATA_IN_FLIGHT + 1, timeout=10)
+        with p2p_lock:
             assert_equal(p.tx_getdata_count, MAX_GETDATA_IN_FLIGHT + 1)
 
         WAIT_TIME = TX_EXPIRY_INTERVAL // 2 + TX_EXPIRY_INTERVAL
         self.log.info("if we wait about {} minutes, we should eventually get more requests".format(WAIT_TIME / 60))
         self.nodes[0].setmocktime(int(time.time() + WAIT_TIME))
-        wait_until(lambda: p.tx_getdata_count == MAX_GETDATA_IN_FLIGHT + 2)
+        p.wait_until(lambda: p.tx_getdata_count == MAX_GETDATA_IN_FLIGHT + 2)
         self.nodes[0].setmocktime(0)
 
     def test_spurious_notfound(self):