[tests] Make p2p-leaktests.py more robust

2 files changed, 8 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()
diff --git a/test/functional/test_framework/mininode.py b/test/functional/test_framework/mininode.py
index d0753276db..5924530230 100755
--- a/test/functional/test_framework/mininode.py
+++ b/test/functional/test_framework/mininode.py
@@ -1837,6 +1837,7 @@ class NetworkThread(Thread):
                     disconnected.append(obj)
             [ obj.handle_close() for obj in disconnected ]
             asyncore.loop(0.1, use_poll=True, map=mininode_socket_map, count=1)
+        logger.debug("Network thread closing")
 
 
 # An exception we can raise if we detect a potential disconnect