Adds helper functions to NodeConnCB

This commit adds some helper functions to NodeConnCB which are useful
for many tests:

- NodeConnCB now keeps track of the number of each message type that
it's received and the most recent message of each type. Many tests
assert on the most recent block, tx or reject message.
- NodeConnCB now keeps track of its connection state by setting a
connected boolean in on_open() and on_close()
- NodeConnCB now has wait_for_block, wait_for_getdata,
wait_for_getheaders, wait_for_inv and wait_for_verack methods

I have updated the individual test cases to make sure that there are no
namespace problems that cause them to fail with these new definitions.
Future commits will remove the duplicate code.

1 files changed, 6 insertions, 2 deletions

diff --git a/test/functional/p2p-leaktests.py b/test/functional/p2p-leaktests.py
index 5853ec86f0..2423d4a969 100755
--- a/test/functional/p2p-leaktests.py
+++ b/test/functional/p2p-leaktests.py
@@ -23,6 +23,7 @@ class CLazyNode(NodeConnCB):
         self.connection = None
         self.unexpected_msg = False
         self.connected = False
+        self.ever_connected = False
 
     def add_connection(self, conn):
         self.connection = conn
@@ -36,6 +37,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)
@@ -121,7 +123,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 +134,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]