1 files changed, 6 insertions, 13 deletions
diff --git a/test/functional/p2p-leaktests.py b/test/functional/p2p-leaktests.py
index 5853ec86f0..4cc3107cf8 100755
--- a/test/functional/p2p-leaktests.py
+++ b/test/functional/p2p-leaktests.py
@@ -20,15 +20,8 @@ banscore = 10
 class CLazyNode(NodeConnCB):
     def __init__(self):
         super().__init__()
-        self.connection = None
         self.unexpected_msg = False
-        self.connected = False
-
-    def add_connection(self, conn):
-        self.connection = conn
-
-    def send_message(self, message):
-        self.connection.send_message(message)
+        self.ever_connected = False
 
     def bad_message(self, message):
         self.unexpected_msg = True
@@ -36,6 +29,7 @@ class CLazyNode(NodeConnCB):
 
     def on_open(self, conn):
         self.connected = True
+        self.ever_connected = True
 
     def on_version(self, conn, message): self.bad_message(message)
     def on_verack(self, conn, message): self.bad_message(message)
@@ -63,9 +57,6 @@ class CLazyNode(NodeConnCB):
 # Node that never sends a version. We'll use this to send a bunch of messages
 # anyway, and eventually get disconnected.
 class CNodeNoVersionBan(CLazyNode):
-    def __init__(self):
-        super().__init__()
-
     # send a bunch of veracks without sending a message. This should get us disconnected.
     # NOTE: implementation-specific check here. Remove if bitcoind ban behavior changes
     def on_open(self, conn):
@@ -121,7 +112,9 @@ class P2PLeakTest(BitcoinTestFramework):
 
         NetworkThread().start()  # Start up network handling in another thread
 
-        assert(wait_until(lambda: no_version_bannode.connected and no_version_idlenode.connected and no_verack_idlenode.version_received, timeout=10))
+        assert wait_until(lambda: no_version_bannode.ever_connected, timeout=10)
+        assert wait_until(lambda: no_version_idlenode.ever_connected, timeout=10)
+        assert wait_until(lambda: no_verack_idlenode.version_received, timeout=10)
 
         # Mine a block and make sure that it's not sent to the connected nodes
         self.nodes[0].generate(1)
@@ -130,7 +123,7 @@ class P2PLeakTest(BitcoinTestFramework):
         time.sleep(5)
 
         #This node should have been banned
-        assert(no_version_bannode.connection.state == "closed")
+        assert not no_version_bannode.connected
 
         [conn.disconnect_node() for conn in connections]