From 2a52ae63bfdde948250df1c876dcdd5af99f03b5 Mon Sep 17 00:00:00 2001
From: John Newbery <john@johnnewbery.com>
Date: Thu, 30 Mar 2017 08:38:46 -0400
Subject: [PATCH] Remove duplicate method definitions in NodeConnCB subclasses

All Node classes in individual test cases subclass from NodeConnCB. Many
have duplicate definitions for methods that are defined in the base
class. This commit removes those duplicate definitions.

This commit removes ~290 lines of duplicate code.
---
 test/functional/maxuploadtarget.py         |  67 ++-------
 test/functional/p2p-acceptblock.py         |  31 +---
 test/functional/p2p-compactblocks.py       | 164 ++++++++-------------
 test/functional/p2p-feefilter.py           |  10 +-
 test/functional/p2p-leaktests.py           |  11 --
 test/functional/p2p-mempool.py             |  53 +------
 test/functional/p2p-segwit.py              |  99 +++----------
 test/functional/p2p-timeouts.py            |  15 +-
 test/functional/p2p-versionbits-warning.py |  18 ---
 test/functional/sendheaders.py             |  99 ++++---------
 test/functional/test_framework/comptool.py |   2 +-
 test/functional/test_framework/mininode.py |  26 +---
 12 files changed, 144 insertions(+), 451 deletions(-)

diff --git a/test/functional/maxuploadtarget.py b/test/functional/maxuploadtarget.py
index 696933512e..b38fdc0f88 100755
--- a/test/functional/maxuploadtarget.py
+++ b/test/functional/maxuploadtarget.py
@@ -10,54 +10,24 @@ if uploadtarget has been reached.
 if uploadtarget has been reached.
 * Verify that the upload counters are reset after 24 hours.
 """
+from collections import defaultdict
+import time
 
 from test_framework.mininode import *
 from test_framework.test_framework import BitcoinTestFramework
 from test_framework.util import *
-import time
 
-# TestNode: bare-bones "peer".  Used mostly as a conduit for a test to sending
-# p2p messages to a node, generating the messages in the main testing logic.
 class TestNode(NodeConnCB):
     def __init__(self):
         super().__init__()
-        self.connection = None
-        self.ping_counter = 1
-        self.last_pong = msg_pong()
-        self.block_receive_map = {}
-
-    def add_connection(self, conn):
-        self.connection = conn
-        self.peer_disconnected = False
+        self.block_receive_map = defaultdict(int)
 
     def on_inv(self, conn, message):
         pass
 
-    # Track the last getdata message we receive (used in the test)
-    def on_getdata(self, conn, message):
-        self.last_getdata = message
-
     def on_block(self, conn, message):
         message.block.calc_sha256()
-        try:
-            self.block_receive_map[message.block.sha256] += 1
-        except KeyError as e:
-            self.block_receive_map[message.block.sha256] = 1
-
-    def wait_for_disconnect(self):
-        def disconnected():
-            return self.peer_disconnected
-        return wait_until(disconnected, timeout=10)
-
-    # Wrapper for the NodeConn's send_message function
-    def send_message(self, message):
-        self.connection.send_message(message)
-
-    def on_pong(self, conn, message):
-        self.last_pong = message
-
-    def on_close(self, conn):
-        self.peer_disconnected = True
+        self.block_receive_map[message.block.sha256] += 1
 
 class MaxUploadTest(BitcoinTestFramework):
  
@@ -183,33 +153,26 @@ class MaxUploadTest(BitcoinTestFramework):
         stop_node(self.nodes[0], 0)
         self.nodes[0] = start_node(0, self.options.tmpdir, ["-whitelist=127.0.0.1", "-maxuploadtarget=1", "-blockmaxsize=999000"])
 
-        #recreate/reconnect 3 test nodes
-        test_nodes = []
-        connections = []
-
-        for i in range(3):
-            test_nodes.append(TestNode())
-            connections.append(NodeConn('127.0.0.1', p2p_port(0), self.nodes[0], test_nodes[i]))
-            test_nodes[i].add_connection(connections[i])
+        #recreate/reconnect a test node
+        test_nodes = [TestNode()]
+        connections = [NodeConn('127.0.0.1', p2p_port(0), self.nodes[0], test_nodes[0])]
+        test_nodes[0].add_connection(connections[0])
 
         NetworkThread().start() # Start up network handling in another thread
-        [x.wait_for_verack() for x in test_nodes]
+        test_nodes[0].wait_for_verack()
 
         #retrieve 20 blocks which should be enough to break the 1MB limit
         getdata_request.inv = [CInv(2, big_new_block)]
         for i in range(20):
-            test_nodes[1].send_message(getdata_request)
-            test_nodes[1].sync_with_ping()
-            assert_equal(test_nodes[1].block_receive_map[big_new_block], i+1)
+            test_nodes[0].send_message(getdata_request)
+            test_nodes[0].sync_with_ping()
+            assert_equal(test_nodes[0].block_receive_map[big_new_block], i+1)
 
         getdata_request.inv = [CInv(2, big_old_block)]
-        test_nodes[1].send_message(getdata_request)
-        test_nodes[1].wait_for_disconnect()
-        assert_equal(len(self.nodes[0].getpeerinfo()), 3) #node is still connected because of the whitelist
+        test_nodes[0].send_and_ping(getdata_request)
+        assert_equal(len(self.nodes[0].getpeerinfo()), 1) #node is still connected because of the whitelist
 
-        self.log.info("Peer 1 still connected after trying to download old block (whitelisted)")
-
-        [c.disconnect_node() for c in connections]
+        self.log.info("Peer still connected after trying to download old block (whitelisted)")
 
 if __name__ == '__main__':
     MaxUploadTest().main()
diff --git a/test/functional/p2p-acceptblock.py b/test/functional/p2p-acceptblock.py
index e2aad3f1d1..01d621a53a 100755
--- a/test/functional/p2p-acceptblock.py
+++ b/test/functional/p2p-acceptblock.py
@@ -54,29 +54,6 @@ from test_framework.util import *
 import time
 from test_framework.blocktools import create_block, create_coinbase
 
-# TestNode: bare-bones "peer".  Used mostly as a conduit for a test to sending
-# p2p messages to a node, generating the messages in the main testing logic.
-class TestNode(NodeConnCB):
-    def __init__(self):
-        super().__init__()
-        self.connection = None
-        self.ping_counter = 1
-        self.last_pong = msg_pong()
-
-    def add_connection(self, conn):
-        self.connection = conn
-
-    # Track the last getdata message we receive (used in the test)
-    def on_getdata(self, conn, message):
-        self.last_getdata = message
-
-    # Wrapper for the NodeConn's send_message function
-    def send_message(self, message):
-        self.connection.send_message(message)
-
-    def on_pong(self, conn, message):
-        self.last_pong = message
-
 class AcceptBlockTest(BitcoinTestFramework):
     def add_options(self, parser):
         parser.add_option("--testbinary", dest="testbinary",
@@ -101,8 +78,8 @@ class AcceptBlockTest(BitcoinTestFramework):
 
     def run_test(self):
         # Setup the p2p connections and start up the network thread.
-        test_node = TestNode()   # connects to node0 (not whitelisted)
-        white_node = TestNode()  # connects to node1 (whitelisted)
+        test_node = NodeConnCB()   # connects to node0 (not whitelisted)
+        white_node = NodeConnCB()  # connects to node1 (whitelisted)
 
         connections = []
         connections.append(NodeConn('127.0.0.1', p2p_port(0), self.nodes[0], test_node))
@@ -227,12 +204,12 @@ class AcceptBlockTest(BitcoinTestFramework):
         # triggers a getdata on block 2 (it should if block 2 is missing).
         with mininode_lock:
             # Clear state so we can check the getdata request
-            test_node.last_getdata = None
+            test_node.last_message.pop("getdata", None)
             test_node.send_message(msg_inv([CInv(2, blocks_h3[0].sha256)]))
 
         test_node.sync_with_ping()
         with mininode_lock:
-            getdata = test_node.last_getdata
+            getdata = test_node.last_message["getdata"]
 
         # Check that the getdata includes the right block
         assert_equal(getdata.inv[0].hash, blocks_h2f[0].sha256)
diff --git a/test/functional/p2p-compactblocks.py b/test/functional/p2p-compactblocks.py
index bf8d113767..cf9608a657 100755
--- a/test/functional/p2p-compactblocks.py
+++ b/test/functional/p2p-compactblocks.py
@@ -19,64 +19,31 @@ class TestNode(NodeConnCB):
     def __init__(self):
         super().__init__()
         self.last_sendcmpct = []
-        self.last_headers = None
-        self.last_inv = None
-        self.last_cmpctblock = None
         self.block_announced = False
-        self.last_getdata = None
-        self.last_getheaders = None
-        self.last_getblocktxn = None
-        self.last_block = None
-        self.last_blocktxn = None
         # Store the hashes of blocks we've seen announced.
         # This is for synchronizing the p2p message traffic,
         # so we can eg wait until a particular block is announced.
-        self.set_announced_blockhashes = set()
-        self.connected = False
-
-    def on_open(self, conn):
-        self.connected = True
-
-    def on_close(self, conn):
-        self.connected = False
+        self.announced_blockhashes = set()
 
     def on_sendcmpct(self, conn, message):
         self.last_sendcmpct.append(message)
 
-    def on_block(self, conn, message):
-        self.last_block = message
-
     def on_cmpctblock(self, conn, message):
-        self.last_cmpctblock = message
         self.block_announced = True
-        self.last_cmpctblock.header_and_shortids.header.calc_sha256()
-        self.set_announced_blockhashes.add(self.last_cmpctblock.header_and_shortids.header.sha256)
+        self.last_message["cmpctblock"].header_and_shortids.header.calc_sha256()
+        self.announced_blockhashes.add(self.last_message["cmpctblock"].header_and_shortids.header.sha256)
 
     def on_headers(self, conn, message):
-        self.last_headers = message
         self.block_announced = True
-        for x in self.last_headers.headers:
+        for x in self.last_message["headers"].headers:
             x.calc_sha256()
-            self.set_announced_blockhashes.add(x.sha256)
+            self.announced_blockhashes.add(x.sha256)
 
     def on_inv(self, conn, message):
-        self.last_inv = message
-        for x in self.last_inv.inv:
+        for x in self.last_message["inv"].inv:
             if x.type == 2:
                 self.block_announced = True
-                self.set_announced_blockhashes.add(x.hash)
-
-    def on_getdata(self, conn, message):
-        self.last_getdata = message
-
-    def on_getheaders(self, conn, message):
-        self.last_getheaders = message
-
-    def on_getblocktxn(self, conn, message):
-        self.last_getblocktxn = message
-
-    def on_blocktxn(self, conn, message):
-        self.last_blocktxn = message
+                self.announced_blockhashes.add(x.hash)
 
     # Requires caller to hold mininode_lock
     def received_block_announcement(self):
@@ -85,9 +52,9 @@ class TestNode(NodeConnCB):
     def clear_block_announcement(self):
         with mininode_lock:
             self.block_announced = False
-            self.last_inv = None
-            self.last_headers = None
-            self.last_cmpctblock = None
+            self.last_message.pop("inv", None)
+            self.last_message.pop("headers", None)
+            self.last_message.pop("cmpctblock", None)
 
     def get_headers(self, locator, hashstop):
         msg = msg_getheaders()
@@ -103,15 +70,14 @@ class TestNode(NodeConnCB):
     def request_headers_and_sync(self, locator, hashstop=0):
         self.clear_block_announcement()
         self.get_headers(locator, hashstop)
-        assert(wait_until(self.received_block_announcement, timeout=30))
-        assert(self.received_block_announcement())
+        assert wait_until(self.received_block_announcement, timeout=30)
         self.clear_block_announcement()
 
     # Block until a block announcement for a particular block hash is
     # received.
     def wait_for_block_announcement(self, block_hash, timeout=30):
         def received_hash():
-            return (block_hash in self.set_announced_blockhashes)
+            return (block_hash in self.announced_blockhashes)
         return wait_until(received_hash, timeout=timeout)
 
     def send_await_disconnect(self, message, timeout=30):
@@ -214,14 +180,14 @@ class CompactBlocksTest(BitcoinTestFramework):
             with mininode_lock:
                 assert predicate(peer), (
                     "block_hash={!r}, cmpctblock={!r}, inv={!r}".format(
-                        block_hash, peer.last_cmpctblock, peer.last_inv))
+                        block_hash, peer.last_message.get("cmpctblock", None), peer.last_message.get("inv", None)))
 
         # We shouldn't get any block announcements via cmpctblock yet.
-        check_announcement_of_new_block(node, test_node, lambda p: p.last_cmpctblock is None)
+        check_announcement_of_new_block(node, test_node, lambda p: "cmpctblock" not in p.last_message)
 
         # Try one more time, this time after requesting headers.
         test_node.request_headers_and_sync(locator=[tip])
-        check_announcement_of_new_block(node, test_node, lambda p: p.last_cmpctblock is None and p.last_inv is not None)
+        check_announcement_of_new_block(node, test_node, lambda p: "cmpctblock" not in p.last_message and "inv" in p.last_message)
 
         # Test a few ways of using sendcmpct that should NOT
         # result in compact block announcements.
@@ -233,7 +199,7 @@ class CompactBlocksTest(BitcoinTestFramework):
         sendcmpct.version = preferred_version+1
         sendcmpct.announce = True
         test_node.send_and_ping(sendcmpct)
-        check_announcement_of_new_block(node, test_node, lambda p: p.last_cmpctblock is None)
+        check_announcement_of_new_block(node, test_node, lambda p: "cmpctblock" not in p.last_message)
 
         # Headers sync before next test.
         test_node.request_headers_and_sync(locator=[tip])
@@ -242,7 +208,7 @@ class CompactBlocksTest(BitcoinTestFramework):
         sendcmpct.version = preferred_version
         sendcmpct.announce = False
         test_node.send_and_ping(sendcmpct)
-        check_announcement_of_new_block(node, test_node, lambda p: p.last_cmpctblock is None)
+        check_announcement_of_new_block(node, test_node, lambda p: "cmpctblock" not in p.last_message)
 
         # Headers sync before next test.
         test_node.request_headers_and_sync(locator=[tip])
@@ -251,26 +217,26 @@ class CompactBlocksTest(BitcoinTestFramework):
         sendcmpct.version = preferred_version
         sendcmpct.announce = True
         test_node.send_and_ping(sendcmpct)
-        check_announcement_of_new_block(node, test_node, lambda p: p.last_cmpctblock is not None)
+        check_announcement_of_new_block(node, test_node, lambda p: "cmpctblock" in p.last_message)
 
         # Try one more time (no headers sync should be needed!)
-        check_announcement_of_new_block(node, test_node, lambda p: p.last_cmpctblock is not None)
+        check_announcement_of_new_block(node, test_node, lambda p: "cmpctblock" in p.last_message)
 
         # Try one more time, after turning on sendheaders
         test_node.send_and_ping(msg_sendheaders())
-        check_announcement_of_new_block(node, test_node, lambda p: p.last_cmpctblock is not None)
+        check_announcement_of_new_block(node, test_node, lambda p: "cmpctblock" in p.last_message)
 
         # Try one more time, after sending a version-1, announce=false message.
         sendcmpct.version = preferred_version-1
         sendcmpct.announce = False
         test_node.send_and_ping(sendcmpct)
-        check_announcement_of_new_block(node, test_node, lambda p: p.last_cmpctblock is not None)
+        check_announcement_of_new_block(node, test_node, lambda p: "cmpctblock" in p.last_message)
 
         # Now turn off announcements
         sendcmpct.version = preferred_version
         sendcmpct.announce = False
         test_node.send_and_ping(sendcmpct)
-        check_announcement_of_new_block(node, test_node, lambda p: p.last_cmpctblock is None and p.last_headers is not None)
+        check_announcement_of_new_block(node, test_node, lambda p: "cmpctblock" not in p.last_message and "headers" in p.last_message)
 
         if old_node is not None:
             # Verify that a peer using an older protocol version can receive
@@ -280,7 +246,7 @@ class CompactBlocksTest(BitcoinTestFramework):
             old_node.send_and_ping(sendcmpct)
             # Header sync
             old_node.request_headers_and_sync(locator=[tip])
-            check_announcement_of_new_block(node, old_node, lambda p: p.last_cmpctblock is not None)
+            check_announcement_of_new_block(node, old_node, lambda p: "cmpctblock" in p.last_message)
 
     # This test actually causes bitcoind to (reasonably!) disconnect us, so do this last.
     def test_invalid_cmpctblock_message(self):
@@ -345,9 +311,9 @@ class CompactBlocksTest(BitcoinTestFramework):
         # Now fetch and check the compact block
         header_and_shortids = None
         with mininode_lock:
-            assert(test_node.last_cmpctblock is not None)
+            assert("cmpctblock" in test_node.last_message)
             # Convert the on-the-wire representation to absolute indexes
-            header_and_shortids = HeaderAndShortIDs(test_node.last_cmpctblock.header_and_shortids)
+            header_and_shortids = HeaderAndShortIDs(test_node.last_message["cmpctblock"].header_and_shortids)
         self.check_compactblock_construction_from_block(version, header_and_shortids, block_hash, block)
 
         # Now fetch the compact block using a normal non-announce getdata
@@ -362,9 +328,9 @@ class CompactBlocksTest(BitcoinTestFramework):
         # Now fetch and check the compact block
         header_and_shortids = None
         with mininode_lock:
-            assert(test_node.last_cmpctblock is not None)
+            assert("cmpctblock" in test_node.last_message)
             # Convert the on-the-wire representation to absolute indexes
-            header_and_shortids = HeaderAndShortIDs(test_node.last_cmpctblock.header_and_shortids)
+            header_and_shortids = HeaderAndShortIDs(test_node.last_message["cmpctblock"].header_and_shortids)
         self.check_compactblock_construction_from_block(version, header_and_shortids, block_hash, block)
 
     def check_compactblock_construction_from_block(self, version, header_and_shortids, block_hash, block):
@@ -424,20 +390,20 @@ class CompactBlocksTest(BitcoinTestFramework):
         for announce in ["inv", "header"]:
             block = self.build_block_on_tip(node, segwit=segwit)
             with mininode_lock:
-                test_node.last_getdata = None
+                test_node.last_message.pop("getdata", None)
 
             if announce == "inv":
                 test_node.send_message(msg_inv([CInv(2, block.sha256)]))
-                success = wait_until(lambda: test_node.last_getheaders is not None, timeout=30)
+                success = wait_until(lambda: "getheaders" in test_node.last_message, timeout=30)
                 assert(success)
                 test_node.send_header_for_blocks([block])
             else:
                 test_node.send_header_for_blocks([block])
-            success = wait_until(lambda: test_node.last_getdata is not None, timeout=30)
+            success = wait_until(lambda: "getdata" in test_node.last_message, timeout=30)
             assert(success)
-            assert_equal(len(test_node.last_getdata.inv), 1)
-            assert_equal(test_node.last_getdata.inv[0].type, 4)
-            assert_equal(test_node.last_getdata.inv[0].hash, block.sha256)
+            assert_equal(len(test_node.last_message["getdata"].inv), 1)
+            assert_equal(test_node.last_message["getdata"].inv[0].type, 4)
+            assert_equal(test_node.last_message["getdata"].inv[0].hash, block.sha256)
 
             # Send back a compactblock message that omits the coinbase
             comp_block = HeaderAndShortIDs()
@@ -453,8 +419,8 @@ class CompactBlocksTest(BitcoinTestFramework):
             assert_equal(int(node.getbestblockhash(), 16), block.hashPrevBlock)
             # Expect a getblocktxn message.
             with mininode_lock:
-                assert(test_node.last_getblocktxn is not None)
-                absolute_indexes = test_node.last_getblocktxn.block_txn_request.to_absolute()
+                assert("getblocktxn" in test_node.last_message)
+                absolute_indexes = test_node.last_message["getblocktxn"].block_txn_request.to_absolute()
             assert_equal(absolute_indexes, [0])  # should be a coinbase request
 
             # Send the coinbase, and verify that the tip advances.
@@ -493,8 +459,8 @@ class CompactBlocksTest(BitcoinTestFramework):
             msg = msg_cmpctblock(compact_block.to_p2p())
             peer.send_and_ping(msg)
             with mininode_lock:
-                assert(peer.last_getblocktxn is not None)
-                absolute_indexes = peer.last_getblocktxn.block_txn_request.to_absolute()
+                assert("getblocktxn" in peer.last_message)
+                absolute_indexes = peer.last_message["getblocktxn"].block_txn_request.to_absolute()
             assert_equal(absolute_indexes, expected_result)
 
         def test_tip_after_message(node, peer, msg, tip):
@@ -558,14 +524,14 @@ class CompactBlocksTest(BitcoinTestFramework):
 
         # Clear out last request.
         with mininode_lock:
-            test_node.last_getblocktxn = None
+            test_node.last_message.pop("getblocktxn", None)
 
         # Send compact block
         comp_block.initialize_from_block(block, prefill_list=[0], use_witness=with_witness)
         test_tip_after_message(node, test_node, msg_cmpctblock(comp_block.to_p2p()), block.sha256)
         with mininode_lock:
             # Shouldn't have gotten a request for any transaction
-            assert(test_node.last_getblocktxn is None)
+            assert("getblocktxn" not in test_node.last_message)
 
     # Incorrectly responding to a getblocktxn shouldn't cause the block to be
     # permanently failed.
@@ -591,8 +557,8 @@ class CompactBlocksTest(BitcoinTestFramework):
         test_node.send_and_ping(msg_cmpctblock(comp_block.to_p2p()))
         absolute_indexes = []
         with mininode_lock:
-            assert(test_node.last_getblocktxn is not None)
-            absolute_indexes = test_node.last_getblocktxn.block_txn_request.to_absolute()
+            assert("getblocktxn" in test_node.last_message)
+            absolute_indexes = test_node.last_message["getblocktxn"].block_txn_request.to_absolute()
         assert_equal(absolute_indexes, [6, 7, 8, 9, 10])
 
         # Now give an incorrect response.
@@ -613,11 +579,11 @@ class CompactBlocksTest(BitcoinTestFramework):
         assert_equal(int(node.getbestblockhash(), 16), block.hashPrevBlock)
 
         # We should receive a getdata request
-        success = wait_until(lambda: test_node.last_getdata is not None, timeout=10)
+        success = wait_until(lambda: "getdata" in test_node.last_message, timeout=10)
         assert(success)
-        assert_equal(len(test_node.last_getdata.inv), 1)
-        assert(test_node.last_getdata.inv[0].type == 2 or test_node.last_getdata.inv[0].type == 2|MSG_WITNESS_FLAG)
-        assert_equal(test_node.last_getdata.inv[0].hash, block.sha256)
+        assert_equal(len(test_node.last_message["getdata"].inv), 1)
+        assert(test_node.last_message["getdata"].inv[0].type == 2 or test_node.last_message["getdata"].inv[0].type == 2|MSG_WITNESS_FLAG)
+        assert_equal(test_node.last_message["getdata"].inv[0].hash, block.sha256)
 
         # Deliver the block
         if version==2:
@@ -641,15 +607,15 @@ class CompactBlocksTest(BitcoinTestFramework):
             num_to_request = random.randint(1, len(block.vtx))
             msg.block_txn_request.from_absolute(sorted(random.sample(range(len(block.vtx)), num_to_request)))
             test_node.send_message(msg)
-            success = wait_until(lambda: test_node.last_blocktxn is not None, timeout=10)
+            success = wait_until(lambda: "blocktxn" in test_node.last_message, timeout=10)
             assert(success)
 
             [tx.calc_sha256() for tx in block.vtx]
             with mininode_lock:
-                assert_equal(test_node.last_blocktxn.block_transactions.blockhash, int(block_hash, 16))
+                assert_equal(test_node.last_message["blocktxn"].block_transactions.blockhash, int(block_hash, 16))
                 all_indices = msg.block_txn_request.to_absolute()
                 for index in all_indices:
-                    tx = test_node.last_blocktxn.block_transactions.transactions.pop(0)
+                    tx = test_node.last_message["blocktxn"].block_transactions.transactions.pop(0)
                     tx.calc_sha256()
                     assert_equal(tx.sha256, block.vtx[index].sha256)
                     if version == 1:
@@ -658,7 +624,7 @@ class CompactBlocksTest(BitcoinTestFramework):
                     else:
                         # Check that the witness matches
                         assert_equal(tx.calc_sha256(True), block.vtx[index].calc_sha256(True))
-                test_node.last_blocktxn = None
+                test_node.last_message.pop("blocktxn", None)
             current_height -= 1
 
         # Next request should send a full block response, as we're past the
@@ -666,13 +632,13 @@ class CompactBlocksTest(BitcoinTestFramework):
         block_hash = node.getblockhash(current_height)
         msg.block_txn_request = BlockTransactionsRequest(int(block_hash, 16), [0])
         with mininode_lock:
-            test_node.last_block = None
-            test_node.last_blocktxn = None
+            test_node.last_message.pop("block", None)
+            test_node.last_message.pop("blocktxn", None)
         test_node.send_and_ping(msg)
         with mininode_lock:
-            test_node.last_block.block.calc_sha256()
-            assert_equal(test_node.last_block.block.sha256, int(block_hash, 16))
-            assert_equal(test_node.last_blocktxn, None)
+            test_node.last_message["block"].block.calc_sha256()
+            assert_equal(test_node.last_message["block"].block.sha256, int(block_hash, 16))
+            assert "blocktxn" not in test_node.last_message
 
     def test_compactblocks_not_at_tip(self, node, test_node):
         # Test that requesting old compactblocks doesn't work.
@@ -685,7 +651,7 @@ class CompactBlocksTest(BitcoinTestFramework):
 
         test_node.clear_block_announcement()
         test_node.send_message(msg_getdata([CInv(4, int(new_blocks[0], 16))]))
-        success = wait_until(lambda: test_node.last_cmpctblock is not None, timeout=30)
+        success = wait_until(lambda: "cmpctblock" in test_node.last_message, timeout=30)
         assert(success)
 
         test_node.clear_block_announcement()
@@ -693,13 +659,13 @@ class CompactBlocksTest(BitcoinTestFramework):
         wait_until(test_node.received_block_announcement, timeout=30)
         test_node.clear_block_announcement()
         with mininode_lock:
-            test_node.last_block = None
+            test_node.last_message.pop("block", None)
         test_node.send_message(msg_getdata([CInv(4, int(new_blocks[0], 16))]))
-        success = wait_until(lambda: test_node.last_block is not None, timeout=30)
+        success = wait_until(lambda: "block" in test_node.last_message, timeout=30)
         assert(success)
         with mininode_lock:
-            test_node.last_block.block.calc_sha256()
-            assert_equal(test_node.last_block.block.sha256, int(new_blocks[0], 16))
+            test_node.last_message["block"].block.calc_sha256()
+            assert_equal(test_node.last_message["block"].block.sha256, int(new_blocks[0], 16))
 
         # Generate an old compactblock, and verify that it's not accepted.
         cur_height = node.getblockcount()
@@ -726,10 +692,10 @@ class CompactBlocksTest(BitcoinTestFramework):
         msg = msg_getblocktxn()
         msg.block_txn_request = BlockTransactionsRequest(block.sha256, [0])
         with mininode_lock:
-            test_node.last_blocktxn = None
+            test_node.last_message.pop("blocktxn", None)
         test_node.send_and_ping(msg)
         with mininode_lock:
-            assert(test_node.last_blocktxn is None)
+            assert "blocktxn" not in test_node.last_message
 
     def activate_segwit(self, node):
         node.generate(144*3)
@@ -750,9 +716,9 @@ class CompactBlocksTest(BitcoinTestFramework):
             wait_until(lambda: l.received_block_announcement(), timeout=30)
         with mininode_lock:
             for l in listeners:
-                assert(l.last_cmpctblock is not None)
-                l.last_cmpctblock.header_and_shortids.header.calc_sha256()
-                assert_equal(l.last_cmpctblock.header_and_shortids.header.sha256, block.sha256)
+                assert "cmpctblock" in l.last_message
+                l.last_message["cmpctblock"].header_and_shortids.header.calc_sha256()
+                assert_equal(l.last_message["cmpctblock"].header_and_shortids.header.sha256, block.sha256)
 
     # Test that we don't get disconnected if we relay a compact block with valid header,
     # but invalid transactions.
@@ -804,7 +770,7 @@ class CompactBlocksTest(BitcoinTestFramework):
             msg = msg_cmpctblock(cmpct_block.to_p2p())
             peer.send_and_ping(msg)
             with mininode_lock:
-                assert(peer.last_getblocktxn is not None)
+                assert "getblocktxn" in peer.last_message
             return block, cmpct_block
 
         block, cmpct_block = announce_cmpct_block(node, stalling_peer)
diff --git a/test/functional/p2p-feefilter.py b/test/functional/p2p-feefilter.py
index 12539be950..fe1a43ac54 100755
--- a/test/functional/p2p-feefilter.py
+++ b/test/functional/p2p-feefilter.py
@@ -22,8 +22,6 @@ def allInvsMatch(invsExpected, testnode):
         time.sleep(1)
     return False
 
-# TestNode: bare-bones "peer".  Used to track which invs are received from a node
-# and to send the node feefilter messages.
 class TestNode(NodeConnCB):
     def __init__(self):
         super().__init__()
@@ -38,10 +36,6 @@ class TestNode(NodeConnCB):
         with mininode_lock:
             self.txinvs = []
 
-    def send_filter(self, feerate):
-        self.send_message(msg_feefilter(feerate))
-        self.sync_with_ping()
-
 class FeeFilterTest(BitcoinTestFramework):
 
     def __init__(self):
@@ -78,7 +72,7 @@ class FeeFilterTest(BitcoinTestFramework):
         test_node.clear_invs()
 
         # Set a filter of 15 sat/byte
-        test_node.send_filter(15000)
+        test_node.send_and_ping(msg_feefilter(15000))
 
         # Test that txs are still being received (paying 20 sat/byte)
         txids = [node1.sendtoaddress(node1.getnewaddress(), 1) for x in range(3)]
@@ -103,7 +97,7 @@ class FeeFilterTest(BitcoinTestFramework):
         test_node.clear_invs()
 
         # Remove fee filter and check that txs are received again
-        test_node.send_filter(0)
+        test_node.send_and_ping(msg_feefilter(0))
         txids = [node1.sendtoaddress(node1.getnewaddress(), 1) for x in range(3)]
         assert(allInvsMatch(txids, test_node))
         test_node.clear_invs()
diff --git a/test/functional/p2p-leaktests.py b/test/functional/p2p-leaktests.py
index 2423d4a969..4cc3107cf8 100755
--- a/test/functional/p2p-leaktests.py
+++ b/test/functional/p2p-leaktests.py
@@ -20,17 +20,9 @@ banscore = 10
 class CLazyNode(NodeConnCB):
     def __init__(self):
         super().__init__()
-        self.connection = None
         self.unexpected_msg = False
-        self.connected = False
         self.ever_connected = False
 
-    def add_connection(self, conn):
-        self.connection = conn
-
-    def send_message(self, message):
-        self.connection.send_message(message)
-
     def bad_message(self, message):
         self.unexpected_msg = True
         self.log.info("should not have received message: %s" % message.command)
@@ -65,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):
diff --git a/test/functional/p2p-mempool.py b/test/functional/p2p-mempool.py
index c17af173c7..188016e718 100755
--- a/test/functional/p2p-mempool.py
+++ b/test/functional/p2p-mempool.py
@@ -12,51 +12,6 @@ from test_framework.mininode import *
 from test_framework.test_framework import BitcoinTestFramework
 from test_framework.util import *
 
-class TestNode(NodeConnCB):
-    def __init__(self):
-        super().__init__()
-        self.connection = None
-        self.ping_counter = 1
-        self.last_pong = msg_pong()
-        self.block_receive_map = {}
-
-    def add_connection(self, conn):
-        self.connection = conn
-        self.peer_disconnected = False
-
-    def on_inv(self, conn, message):
-        pass
-
-    # Track the last getdata message we receive (used in the test)
-    def on_getdata(self, conn, message):
-        self.last_getdata = message
-
-    def on_block(self, conn, message):
-        message.block.calc_sha256()
-        try:
-            self.block_receive_map[message.block.sha256] += 1
-        except KeyError as e:
-            self.block_receive_map[message.block.sha256] = 1
-
-    def wait_for_disconnect(self):
-        def disconnected():
-            return self.peer_disconnected
-        return wait_until(disconnected, timeout=10)
-
-    # Wrapper for the NodeConn's send_message function
-    def send_message(self, message):
-        self.connection.send_message(message)
-
-    def on_pong(self, conn, message):
-        self.last_pong = message
-
-    def on_close(self, conn):
-        self.peer_disconnected = True
-
-    def send_mempool(self):
-        self.lastInv = []
-        self.send_message(msg_mempool())
-
 class P2PMempoolTests(BitcoinTestFramework):
 
     def __init__(self):
@@ -65,20 +20,18 @@ class P2PMempoolTests(BitcoinTestFramework):
         self.num_nodes = 2
 
     def setup_network(self):
-        # Start a node with maxuploadtarget of 200 MB (/24h)
-        self.nodes = []
-        self.nodes.append(start_node(0, self.options.tmpdir, ["-peerbloomfilters=0"]))
+        self.nodes = [start_node(0, self.options.tmpdir, ["-peerbloomfilters=0"])]
 
     def run_test(self):
         #connect a mininode
-        aTestNode = TestNode()
+        aTestNode = NodeConnCB()
         node = NodeConn('127.0.0.1', p2p_port(0), self.nodes[0], aTestNode)
         aTestNode.add_connection(node)
         NetworkThread().start()
         aTestNode.wait_for_verack()
 
         #request mempool
-        aTestNode.send_mempool()
+        aTestNode.send_message(msg_mempool())
         aTestNode.wait_for_disconnect()
 
         #mininode must be disconnected at this point
diff --git a/test/functional/p2p-segwit.py b/test/functional/p2p-segwit.py
index cd7b788eb4..0ed5fd3eaf 100755
--- a/test/functional/p2p-segwit.py
+++ b/test/functional/p2p-segwit.py
@@ -35,79 +35,22 @@ def get_virtual_size(witness_block):
 class TestNode(NodeConnCB):
     def __init__(self):
         super().__init__()
-        self.connection = None
-        self.ping_counter = 1
-        self.last_pong = msg_pong(0)
-        self.sleep_time = 0.05
         self.getdataset = set()
-        self.last_reject = None
-
-    def add_connection(self, conn):
-        self.connection = conn
-
-    # Wrapper for the NodeConn's send_message function
-    def send_message(self, message):
-        self.connection.send_message(message)
-
-    def on_inv(self, conn, message):
-        self.last_inv = message
-
-    def on_block(self, conn, message):
-        self.last_block = message.block
-        self.last_block.calc_sha256()
 
     def on_getdata(self, conn, message):
         for inv in message.inv:
             self.getdataset.add(inv.hash)
-        self.last_getdata = message
-
-    def on_getheaders(self, conn, message):
-        self.last_getheaders = message
-
-    def on_pong(self, conn, message):
-        self.last_pong = message
-
-    def on_reject(self, conn, message):
-        self.last_reject = message
-
-    # Syncing helpers
-    def sync(self, test_function, timeout=60):
-        while timeout > 0:
-            with mininode_lock:
-                if test_function():
-                    return
-            time.sleep(self.sleep_time)
-            timeout -= self.sleep_time
-        raise AssertionError("Sync failed to complete")
-        
-    def wait_for_block(self, blockhash, timeout=60):
-        test_function = lambda: self.last_block != None and self.last_block.sha256 == blockhash
-        self.sync(test_function, timeout)
-        return
-
-    def wait_for_getdata(self, timeout=60):
-        test_function = lambda: self.last_getdata != None
-        self.sync(test_function, timeout)
-
-    def wait_for_getheaders(self, timeout=60):
-        test_function = lambda: self.last_getheaders != None
-        self.sync(test_function, timeout)
-
-    def wait_for_inv(self, expected_inv, timeout=60):
-        test_function = lambda: self.last_inv != expected_inv
-        self.sync(test_function, timeout)
 
     def announce_tx_and_wait_for_getdata(self, tx, timeout=60):
         with mininode_lock:
-            self.last_getdata = None
+            self.last_message.pop("getdata", None)
         self.send_message(msg_inv(inv=[CInv(1, tx.sha256)]))
         self.wait_for_getdata(timeout)
-        return
 
     def announce_block_and_wait_for_getdata(self, block, use_header, timeout=60):
         with mininode_lock:
-            self.last_getdata = None
-            self.last_getheaders = None
+            self.last_message.pop("getdata", None)
+            self.last_message.pop("getheaders", None)
         msg = msg_headers()
         msg.headers = [ CBlockHeader(block) ]
         if use_header:
@@ -117,11 +60,10 @@ class TestNode(NodeConnCB):
             self.wait_for_getheaders()
             self.send_message(msg)
         self.wait_for_getdata()
-        return
 
     def announce_block(self, block, use_header):
         with mininode_lock:
-            self.last_getdata = None
+            self.last_message.pop("getdata", None)
         if use_header:
             msg = msg_headers()
             msg.headers = [ CBlockHeader(block) ]
@@ -131,22 +73,22 @@ class TestNode(NodeConnCB):
 
     def request_block(self, blockhash, inv_type, timeout=60):
         with mininode_lock:
-            self.last_block = None
+            self.last_message.pop("block", None)
         self.send_message(msg_getdata(inv=[CInv(inv_type, blockhash)]))
         self.wait_for_block(blockhash, timeout)
-        return self.last_block
+        return self.last_message["block"].block
 
     def test_transaction_acceptance(self, tx, with_witness, accepted, reason=None):
         tx_message = msg_tx(tx)
         if with_witness:
             tx_message = msg_witness_tx(tx)
         self.send_message(tx_message)
-        self.sync_with_ping(60)
+        self.sync_with_ping()
         assert_equal(tx.hash in self.connection.rpc.getrawmempool(), accepted)
         if (reason != None and not accepted):
             # Check the rejection reason as well.
             with mininode_lock:
-                assert_equal(self.last_reject.reason, reason)
+                assert_equal(self.last_message["reject"].reason, reason)
 
     # Test whether a witness block had the correct effect on the tip
     def test_witness_block(self, block, accepted, with_witness=True):
@@ -154,10 +96,9 @@ class TestNode(NodeConnCB):
             self.send_message(msg_witness_block(block))
         else:
             self.send_message(msg_block(block))
-        self.sync_with_ping(60)
+        self.sync_with_ping()
         assert_equal(self.connection.rpc.getbestblockhash() == block.hash, accepted)
 
-
 # Used to keep track of anyone-can-spend outputs that we can use in the tests
 class UTXO(object):
     def __init__(self, sha256, n, nValue):
@@ -228,7 +169,7 @@ class SegWitTest(BitcoinTestFramework):
         block = self.build_next_block(nVersion=1)
         block.solve()
         self.test_node.send_message(msg_block(block))
-        self.test_node.sync_with_ping(60) # make sure the block was processed
+        self.test_node.sync_with_ping() # make sure the block was processed
         txid = block.vtx[0].sha256
 
         self.nodes[0].generate(99) # let the block mature
@@ -244,7 +185,7 @@ class SegWitTest(BitcoinTestFramework):
         assert_equal(msg_tx(tx).serialize(), msg_witness_tx(tx).serialize())
 
         self.test_node.send_message(msg_witness_tx(tx))
-        self.test_node.sync_with_ping(60) # make sure the tx was processed
+        self.test_node.sync_with_ping() # make sure the tx was processed
         assert(tx.hash in self.nodes[0].getrawmempool())
         # Save this transaction for later
         self.utxo.append(UTXO(tx.sha256, 0, 49*100000000))
@@ -279,12 +220,12 @@ class SegWitTest(BitcoinTestFramework):
         # TODO: fix synchronization so we can test reject reason
         # Right now, bitcoind delays sending reject messages for blocks
         # until the future, making synchronization here difficult.
-        #assert_equal(self.test_node.last_reject.reason, "unexpected-witness")
+        #assert_equal(self.test_node.last_message["reject"].reason, "unexpected-witness")
 
         # But it should not be permanently marked bad...
         # Resend without witness information.
         self.test_node.send_message(msg_block(block))
-        self.test_node.sync_with_ping(60)
+        self.test_node.sync_with_ping()
         assert_equal(self.nodes[0].getbestblockhash(), block.hash)
 
         sync_blocks(self.nodes)
@@ -893,7 +834,7 @@ class SegWitTest(BitcoinTestFramework):
         # Verify that if a peer doesn't set nServices to include NODE_WITNESS,
         # the getdata is just for the non-witness portion.
         self.old_node.announce_tx_and_wait_for_getdata(tx)
-        assert(self.old_node.last_getdata.inv[0].type == 1)
+        assert(self.old_node.last_message["getdata"].inv[0].type == 1)
 
         # Since we haven't delivered the tx yet, inv'ing the same tx from
         # a witness transaction ought not result in a getdata.
@@ -1028,20 +969,20 @@ class SegWitTest(BitcoinTestFramework):
         block1.solve()
 
         self.test_node.announce_block_and_wait_for_getdata(block1, use_header=False)
-        assert(self.test_node.last_getdata.inv[0].type == blocktype)
+        assert(self.test_node.last_message["getdata"].inv[0].type == blocktype)
         self.test_node.test_witness_block(block1, True)
 
         block2 = self.build_next_block(nVersion=4)
         block2.solve()
 
         self.test_node.announce_block_and_wait_for_getdata(block2, use_header=True)
-        assert(self.test_node.last_getdata.inv[0].type == blocktype)
+        assert(self.test_node.last_message["getdata"].inv[0].type == blocktype)
         self.test_node.test_witness_block(block2, True)
 
         block3 = self.build_next_block(nVersion=(VB_TOP_BITS | (1<<15)))
         block3.solve()
         self.test_node.announce_block_and_wait_for_getdata(block3, use_header=True)
-        assert(self.test_node.last_getdata.inv[0].type == blocktype)
+        assert(self.test_node.last_message["getdata"].inv[0].type == blocktype)
         self.test_node.test_witness_block(block3, True)
 
         # Check that we can getdata for witness blocks or regular blocks,
@@ -1250,9 +1191,9 @@ class SegWitTest(BitcoinTestFramework):
         # Spending a higher version witness output is not allowed by policy,
         # even with fRequireStandard=false.
         self.test_node.test_transaction_acceptance(tx3, with_witness=True, accepted=False)
-        self.test_node.sync_with_ping(60)
+        self.test_node.sync_with_ping()
         with mininode_lock:
-            assert(b"reserved for soft-fork upgrades" in self.test_node.last_reject.reason)
+            assert(b"reserved for soft-fork upgrades" in self.test_node.last_message["reject"].reason)
 
         # Building a block with the transaction must be valid, however.
         block = self.build_next_block()
@@ -1380,7 +1321,7 @@ class SegWitTest(BitcoinTestFramework):
         for i in range(NUM_TESTS):
             # Ping regularly to keep the connection alive
             if (not i % 100):
-                self.test_node.sync_with_ping(60)
+                self.test_node.sync_with_ping()
             # Choose random number of inputs to use.
             num_inputs = random.randint(1, 10)
             # Create a slight bias for producing more utxos
diff --git a/test/functional/p2p-timeouts.py b/test/functional/p2p-timeouts.py
index de4edd6800..210e561e74 100755
--- a/test/functional/p2p-timeouts.py
+++ b/test/functional/p2p-timeouts.py
@@ -28,20 +28,9 @@ from test_framework.test_framework import BitcoinTestFramework
 from test_framework.util import *
 
 class TestNode(NodeConnCB):
-    def __init__(self):
-        super().__init__()
-        self.connected = False
-        self.received_version = False
-
-    def on_open(self, conn):
-        self.connected = True
-
-    def on_close(self, conn):
-        self.connected = False
-
     def on_version(self, conn, message):
         # Don't send a verack in response
-        self.received_version = True
+        pass
 
 class TimeoutsTest(BitcoinTestFramework):
     def __init__(self):
@@ -83,7 +72,7 @@ class TimeoutsTest(BitcoinTestFramework):
 
         sleep(30)
 
-        assert(self.no_verack_node.received_version)
+        assert "version" in self.no_verack_node.last_message
 
         assert(self.no_verack_node.connected)
         assert(self.no_version_node.connected)
diff --git a/test/functional/p2p-versionbits-warning.py b/test/functional/p2p-versionbits-warning.py
index da960e2d80..06cdd5f91d 100755
--- a/test/functional/p2p-versionbits-warning.py
+++ b/test/functional/p2p-versionbits-warning.py
@@ -24,28 +24,10 @@ WARN_UNKNOWN_RULES_MINED = "Unknown block versions being mined! It's possible un
 WARN_UNKNOWN_RULES_ACTIVE = "unknown new rules activated (versionbit {})".format(VB_UNKNOWN_BIT)
 VB_PATTERN = re.compile("^Warning.*versionbit")
 
-# TestNode: bare-bones "peer".  Used mostly as a conduit for a test to sending
-# p2p messages to a node, generating the messages in the main testing logic.
 class TestNode(NodeConnCB):
-    def __init__(self):
-        super().__init__()
-        self.connection = None
-        self.ping_counter = 1
-        self.last_pong = msg_pong()
-
-    def add_connection(self, conn):
-        self.connection = conn
-
     def on_inv(self, conn, message):
         pass
 
-    # Wrapper for the NodeConn's send_message function
-    def send_message(self, message):
-        self.connection.send_message(message)
-
-    def on_pong(self, conn, message):
-        self.last_pong = message
-
 class VersionBitsWarningTest(BitcoinTestFramework):
     def __init__(self):
         super().__init__()
diff --git a/test/functional/sendheaders.py b/test/functional/sendheaders.py
index 1a7475ae84..ef20d42238 100755
--- a/test/functional/sendheaders.py
+++ b/test/functional/sendheaders.py
@@ -81,23 +81,17 @@ from test_framework.blocktools import create_block, create_coinbase
 
 direct_fetch_response_time = 0.05
 
-class BaseNode(NodeConnCB):
+class TestNode(NodeConnCB):
     def __init__(self):
         super().__init__()
-        self.last_inv = None
-        self.last_headers = None
-        self.last_block = None
-        self.last_getdata = None
         self.block_announced = False
-        self.last_getheaders = None
-        self.disconnected = False
         self.last_blockhash_announced = None
 
     def clear_last_announcement(self):
         with mininode_lock:
             self.block_announced = False
-            self.last_inv = None
-            self.last_headers = None
+            self.last_message.pop("inv", None)
+            self.last_message.pop("headers", None)
 
     # Request data for a list of block hashes
     def get_data(self, block_hashes):
@@ -118,29 +112,17 @@ class BaseNode(NodeConnCB):
         self.connection.send_message(msg)
 
     def on_inv(self, conn, message):
-        self.last_inv = message
         self.block_announced = True
         self.last_blockhash_announced = message.inv[-1].hash
 
     def on_headers(self, conn, message):
-        self.last_headers = message
         if len(message.headers):
             self.block_announced = True
             message.headers[-1].calc_sha256()
             self.last_blockhash_announced = message.headers[-1].sha256
 
     def on_block(self, conn, message):
-        self.last_block = message.block
-        self.last_block.calc_sha256()
-
-    def on_getdata(self, conn, message):
-        self.last_getdata = message
-
-    def on_getheaders(self, conn, message):
-        self.last_getheaders = message
-
-    def on_close(self, conn):
-        self.disconnected = True
+        self.last_message["block"].calc_sha256()
 
     # Test whether the last announcement we received had the
     # right header or the right inv
@@ -155,43 +137,27 @@ class BaseNode(NodeConnCB):
 
             success = True
             compare_inv = []
-            if self.last_inv != None:
-                compare_inv = [x.hash for x in self.last_inv.inv]
+            if "inv" in self.last_message:
+                compare_inv = [x.hash for x in self.last_message["inv"].inv]
             if compare_inv != expect_inv:
                 success = False
 
             hash_headers = []
-            if self.last_headers != None:
+            if "headers" in self.last_message:
                 # treat headers as a list of block hashes
-                hash_headers = [ x.sha256 for x in self.last_headers.headers ]
+                hash_headers = [ x.sha256 for x in self.last_message["headers"].headers ]
             if hash_headers != expect_headers:
                 success = False
 
-            self.last_inv = None
-            self.last_headers = None
+            self.last_message.pop("inv", None)
+            self.last_message.pop("headers", None)
         return success
 
-    # Syncing helpers
-    def wait_for_block(self, blockhash, timeout=60):
-        test_function = lambda: self.last_block != None and self.last_block.sha256 == blockhash
-        assert(wait_until(test_function, timeout=timeout))
-        return
-
-    def wait_for_getheaders(self, timeout=60):
-        test_function = lambda: self.last_getheaders != None
-        assert(wait_until(test_function, timeout=timeout))
-        return
-
     def wait_for_getdata(self, hash_list, timeout=60):
         if hash_list == []:
             return
 
-        test_function = lambda: self.last_getdata != None and [x.hash for x in self.last_getdata.inv] == hash_list
-        assert(wait_until(test_function, timeout=timeout))
-        return
-
-    def wait_for_disconnect(self, timeout=60):
-        test_function = lambda: self.disconnected
+        test_function = lambda: "getdata" in self.last_message and [x.hash for x in self.last_message["getdata"].inv] == hash_list
         assert(wait_until(test_function, timeout=timeout))
         return
 
@@ -210,17 +176,6 @@ class BaseNode(NodeConnCB):
         getblocks_message.locator.vHave = locator
         self.send_message(getblocks_message)
 
-# InvNode: This peer should only ever receive inv's, because it doesn't ever send a
-# "sendheaders" message.
-class InvNode(BaseNode):
-    def __init__(self):
-        BaseNode.__init__(self)
-
-# TestNode: This peer is the one we use for most of the testing.
-class TestNode(BaseNode):
-    def __init__(self):
-        BaseNode.__init__(self)
-
 class SendHeadersTest(BitcoinTestFramework):
     def __init__(self):
         super().__init__()
@@ -260,7 +215,7 @@ class SendHeadersTest(BitcoinTestFramework):
 
     def run_test(self):
         # Setup the p2p connections and start up the network thread.
-        inv_node = InvNode()
+        inv_node = TestNode()
         test_node = TestNode()
 
         self.p2p_connections = [inv_node, test_node]
@@ -368,8 +323,8 @@ class SendHeadersTest(BitcoinTestFramework):
                 inv_node.sync_with_ping()
                 # This block should not be announced to the inv node (since it also
                 # broadcast it)
-                assert_equal(inv_node.last_inv, None)
-                assert_equal(inv_node.last_headers, None)
+                assert "inv" not in inv_node.last_message
+                assert "headers" not in inv_node.last_message
                 tip = self.mine_blocks(1)
                 assert_equal(inv_node.check_last_announcement(inv=[tip]), True)
                 assert_equal(test_node.check_last_announcement(headers=[tip]), True)
@@ -459,12 +414,12 @@ class SendHeadersTest(BitcoinTestFramework):
             inv_node.send_message(msg_block(blocks[-1]))
 
         inv_node.sync_with_ping() # Make sure blocks are processed
-        test_node.last_getdata = None
+        test_node.last_message.pop("getdata", None)
         test_node.send_header_for_blocks(blocks)
         test_node.sync_with_ping()
         # should not have received any getdata messages
         with mininode_lock:
-            assert_equal(test_node.last_getdata, None)
+            assert "getdata" not in test_node.last_message
 
         # This time, direct fetch should work
         blocks = []
@@ -498,11 +453,11 @@ class SendHeadersTest(BitcoinTestFramework):
 
         # Announcing one block on fork should not trigger direct fetch
         # (less work than tip)
-        test_node.last_getdata = None
+        test_node.last_message.pop("getdata", None)
         test_node.send_header_for_blocks(blocks[0:1])
         test_node.sync_with_ping()
         with mininode_lock:
-            assert_equal(test_node.last_getdata, None)
+            assert "getdata" not in test_node.last_message
 
         # Announcing one more block on fork should trigger direct fetch for
         # both blocks (same work as tip)
@@ -517,11 +472,11 @@ class SendHeadersTest(BitcoinTestFramework):
         test_node.wait_for_getdata([x.sha256 for x in blocks[2:16]], timeout=direct_fetch_response_time)
 
         # Announcing 1 more header should not trigger any response
-        test_node.last_getdata = None
+        test_node.last_message.pop("getdata", None)
         test_node.send_header_for_blocks(blocks[18:19])
         test_node.sync_with_ping()
         with mininode_lock:
-            assert_equal(test_node.last_getdata, None)
+            assert "getdata" not in test_node.last_message
 
         self.log.info("Part 4: success!")
 
@@ -532,7 +487,7 @@ class SendHeadersTest(BitcoinTestFramework):
         # First we test that receipt of an unconnecting header doesn't prevent
         # chain sync.
         for i in range(10):
-            test_node.last_getdata = None
+            test_node.last_message.pop("getdata", None)
             blocks = []
             # Create two more blocks.
             for j in range(2):
@@ -543,7 +498,7 @@ class SendHeadersTest(BitcoinTestFramework):
                 height += 1
             # Send the header of the second block -> this won't connect.
             with mininode_lock:
-                test_node.last_getheaders = None
+                test_node.last_message.pop("getheaders", None)
             test_node.send_header_for_blocks([blocks[1]])
             test_node.wait_for_getheaders(timeout=1)
             test_node.send_header_for_blocks(blocks)
@@ -566,7 +521,7 @@ class SendHeadersTest(BitcoinTestFramework):
         for i in range(1, MAX_UNCONNECTING_HEADERS):
             # Send a header that doesn't connect, check that we get a getheaders.
             with mininode_lock:
-                test_node.last_getheaders = None
+                test_node.last_message.pop("getheaders", None)
             test_node.send_header_for_blocks([blocks[i]])
             test_node.wait_for_getheaders(timeout=1)
 
@@ -581,25 +536,21 @@ class SendHeadersTest(BitcoinTestFramework):
         for i in range(5*MAX_UNCONNECTING_HEADERS - 1):
             # Send a header that doesn't connect, check that we get a getheaders.
             with mininode_lock:
-                test_node.last_getheaders = None
+                test_node.last_message.pop("getheaders", None)
             test_node.send_header_for_blocks([blocks[i%len(blocks)]])
             test_node.wait_for_getheaders(timeout=1)
 
         # Eventually this stops working.
-        with mininode_lock:
-            self.last_getheaders = None
         test_node.send_header_for_blocks([blocks[-1]])
 
         # Should get disconnected
         test_node.wait_for_disconnect()
-        with mininode_lock:
-            self.last_getheaders = True
 
         self.log.info("Part 5: success!")
 
         # Finally, check that the inv node never received a getdata request,
         # throughout the test
-        assert_equal(inv_node.last_getdata, None)
+        assert "getdata" not in inv_node.last_message
 
 if __name__ == '__main__':
     SendHeadersTest().main()
diff --git a/test/functional/test_framework/comptool.py b/test/functional/test_framework/comptool.py
index cb9a938f94..9f062865a3 100755
--- a/test/functional/test_framework/comptool.py
+++ b/test/functional/test_framework/comptool.py
@@ -192,7 +192,7 @@ class TestManager(object):
         return wait_until(disconnected, timeout=10)
 
     def wait_for_verack(self):
-        [node.wait_for_verack() for node in self.test_nodes]
+        return all(node.wait_for_verack() for node in self.test_nodes)
 
     def wait_for_pings(self, counter):
         def received_pongs():
diff --git a/test/functional/test_framework/mininode.py b/test/functional/test_framework/mininode.py
index 3b4d05df6c..b2531a6c9e 100755
--- a/test/functional/test_framework/mininode.py
+++ b/test/functional/test_framework/mininode.py
@@ -1590,34 +1590,25 @@ class NodeConnCB(object):
 
     # Message receiving helper methods
 
-    def sync(self, test_function, timeout=60):
-        while timeout > 0:
-            with mininode_lock:
-                if test_function():
-                    return
-            time.sleep(0.05)
-            timeout -= 0.05
-        raise AssertionError("Sync failed to complete")
-
     def wait_for_block(self, blockhash, timeout=60):
         test_function = lambda: self.last_message.get("block") and self.last_message["block"].block.rehash() == blockhash
-        self.sync(test_function, timeout)
+        assert wait_until(test_function, timeout=timeout)
 
     def wait_for_getdata(self, timeout=60):
         test_function = lambda: self.last_message.get("getdata")
-        self.sync(test_function, timeout)
+        assert wait_until(test_function, timeout=timeout)
 
     def wait_for_getheaders(self, timeout=60):
         test_function = lambda: self.last_message.get("getheaders")
-        self.sync(test_function, timeout)
+        assert wait_until(test_function, timeout=timeout)
 
     def wait_for_inv(self, expected_inv, timeout=60):
         test_function = lambda: self.last_message.get("inv") and self.last_message["inv"] != expected_inv
-        self.sync(test_function, timeout)
+        assert wait_until(test_function, timeout=timeout)
 
     def wait_for_verack(self, timeout=60):
         test_function = lambda: self.message_count["verack"]
-        self.sync(test_function, timeout)
+        assert wait_until(test_function, timeout=timeout)
 
     # Message sending helper functions
 
@@ -1635,12 +1626,9 @@ class NodeConnCB(object):
     def sync_with_ping(self, timeout=60):
         self.send_message(msg_ping(nonce=self.ping_counter))
         test_function = lambda: self.last_message.get("pong") and self.last_message["pong"].nonce == self.ping_counter
-        success = wait_until(test_function, timeout = timeout)
-        if not success:
-            logger.error("sync_with_ping failed!")
-            raise AssertionError("sync_with_ping failed!")
+        assert wait_until(test_function, timeout=timeout)
         self.ping_counter += 1
-        return success
+        return True
 
 # The actual NodeConn class
 # This class provides an interface for a p2p connection to a specified node