[tests] Make p2p-leaktests.py more robust

1 files changed, 7 insertions, 2 deletions

diff --git a/test/functional/p2p-leaktests.py b/test/functional/p2p-leaktests.py
index f0d4d9a8b8..26f6f86aff 100755
--- a/test/functional/p2p-leaktests.py
+++ b/test/functional/p2p-leaktests.py
@@ -140,6 +140,9 @@ class P2PLeakTest(BitcoinTestFramework):
 
         [conn.disconnect_node() for conn in connections]
 
+        # Wait until all connections are closed
+        wait_until(lambda: len(self.nodes[0].getpeerinfo()) == 0)
+
         # Make sure no unexpected messages came in
         assert(no_version_bannode.unexpected_msg == False)
         assert(no_version_idlenode.unexpected_msg == False)
@@ -158,8 +161,10 @@ class P2PLeakTest(BitcoinTestFramework):
         allowed_service_bit5_node.add_connection(connections[5])
         allowed_service_bit7_node.add_connection(connections[6])
 
-        wait_until(lambda: allowed_service_bit5_node.message_count["verack"], timeout=10, lock=mininode_lock)
-        wait_until(lambda: allowed_service_bit7_node.message_count["verack"], timeout=10, lock=mininode_lock)
+        NetworkThread().start()  # Network thread stopped when all previous NodeConnCBs disconnected. Restart it
+
+        wait_until(lambda: allowed_service_bit5_node.message_count["verack"], lock=mininode_lock)
+        wait_until(lambda: allowed_service_bit7_node.message_count["verack"], lock=mininode_lock)
 
 if __name__ == '__main__':
     P2PLeakTest().main()