Merge #10169: [tests] Remove func test code duplication

2a52ae6 Remove duplicate method definitions in NodeConnCB subclasses (John Newbery)
52e15aa Adds helper functions to NodeConnCB (John Newbery)

Tree-SHA512: 2d7909eb85b3bde0fc3ebf133798eca21e561f4b2a2880937750820a42856cfb61fc94e30591c14ac13218bcfae0ebe7c5e8662a7b10f5b02470325c44a86cf1

test/functional/maxuploadtarget.py

@@ -10,63 +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
-
-    # Spin until verack message is received from the node.
-    # We use this to signal that our test can begin. This
-    # is called from the testing thread, so it needs to acquire
-    # the global lock.
-    def wait_for_verack(self):
-        def veracked():
-            return self.verack_received
-        return wait_until(veracked, timeout=10)
-
-    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):
  
@@ -192,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()

test/functional/p2p-acceptblock.py

@@ -54,40 +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
-
-    # Spin until verack message is received from the node.
-    # We use this to signal that our test can begin. This
-    # is called from the testing thread, so it needs to acquire
-    # the global lock.
-    def wait_for_verack(self):
-        while True:
-            with mininode_lock:
-                if self.verack_received:
-                    return
-            time.sleep(0.05)
-
-    # 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",
@@ -112,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))
@@ -238,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)

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)

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()

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]
 

test/functional/p2p-mempool.py

@@ -12,60 +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
-
-    # Spin until verack message is received from the node.
-    # We use this to signal that our test can begin. This
-    # is called from the testing thread, so it needs to acquire
-    # the global lock.
-    def wait_for_verack(self):
-        def veracked():
-            return self.verack_received
-        return wait_until(veracked, timeout=10)
-
-    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):
@@ -74,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

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

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)

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__()

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()

test/functional/test_framework/comptool.py

@@ -192,9 +192,7 @@ class TestManager(object):
         return wait_until(disconnected, timeout=10)
 
     def wait_for_verack(self):
-        def veracked():
-            return all(node.verack_received for node in self.test_nodes)
-        return wait_until(veracked, timeout=10)
+        return all(node.wait_for_verack() for node in self.test_nodes)
 
     def wait_for_pings(self, counter):
         def received_pongs():

test/functional/test_framework/mininode.py

@@ -20,21 +20,22 @@ msg_block, msg_tx, msg_headers, etc.:
 ser_*, deser_*: functions that handle serialization/deserialization
 """
 
-import struct
-import socket
 import asyncore
-import time
-import sys
-import random
-from .util import hex_str_to_bytes, bytes_to_hex_str
-from io import BytesIO
 from codecs import encode
+from collections import defaultdict
+import copy
 import hashlib
-from threading import RLock
-from threading import Thread
+from io import BytesIO
 import logging
-import copy
+import random
+import socket
+import struct
+import sys
+import time
+from threading import RLock, Thread
+
 from test_framework.siphash import siphash256
+from test_framework.util import hex_str_to_bytes, bytes_to_hex_str
 
 BIP0031_VERSION = 60000
 MY_VERSION = 70014  # past bip-31 for ping/pong
@@ -1466,30 +1467,57 @@ class msg_witness_blocktxn(msg_blocktxn):
         r += self.block_transactions.serialize(with_witness=True)
         return r
 
-# This is what a callback should look like for NodeConn
-# Reimplement the on_* functions to provide handling for events
 class NodeConnCB(object):
+    """Callback and helper functions for P2P connection to a bitcoind node.
+
+    Individual testcases should subclass this and override the on_* methods
+    if they want to alter message handling behaviour.
+    """
+
     def __init__(self):
-        self.verack_received = False
+        # Track whether we have a P2P connection open to the node
+        self.connected = False
+        self.connection = None
+
+        # Track number of messages of each type received and the most recent
+        # message of each type
+        self.message_count = defaultdict(int)
+        self.last_message = {}
+
+        # A count of the number of ping messages we've sent to the node
+        self.ping_counter = 1
+
         # deliver_sleep_time is helpful for debugging race conditions in p2p
         # tests; it causes message delivery to sleep for the specified time
         # before acquiring the global lock and delivering the next message.
         self.deliver_sleep_time = None
+
         # Remember the services our peer has advertised
         self.peer_services = None
-        self.connection = None
-        self.ping_counter = 1
-        self.last_pong = msg_pong()
+
+    # Message receiving methods
 
     def deliver(self, conn, message):
+        """Receive message and dispatch message to appropriate callback.
+
+        We keep a count of how many of each message type has been received
+        and the most recent message of each type.
+
+        Optionally waits for deliver_sleep_time before dispatching message.
+        """
+
         deliver_sleep = self.get_deliver_sleep_time()
         if deliver_sleep is not None:
             time.sleep(deliver_sleep)
         with mininode_lock:
             try:
-                getattr(self, 'on_' + message.command.decode('ascii'))(conn, message)
+                command = message.command.decode('ascii')
+                self.message_count[command] += 1
+                self.last_message[command] = message
+                getattr(self, 'on_' + command)(conn, message)
             except:
-                logger.exception("ERROR delivering %s" % repr(message))
+                print("ERROR delivering %s (%s)" % (repr(message),
+                                                    sys.exc_info()[0]))
 
     def set_deliver_sleep_time(self, value):
         with mininode_lock:
@@ -1499,14 +1527,20 @@ class NodeConnCB(object):
         with mininode_lock:
             return self.deliver_sleep_time
 
-    # Callbacks which can be overridden by subclasses
-    #################################################
+    # Callback methods. Can be overridden by subclasses in individual test
+    # cases to provide custom message handling behaviour.
+
+    def on_open(self, conn):
+        self.connected = True
+
+    def on_close(self, conn):
+        self.connected = False
+        self.connection = None
 
     def on_addr(self, conn, message): pass
     def on_alert(self, conn, message): pass
     def on_block(self, conn, message): pass
     def on_blocktxn(self, conn, message): pass
-    def on_close(self, conn): pass
     def on_cmpctblock(self, conn, message): pass
     def on_feefilter(self, conn, message): pass
     def on_getaddr(self, conn, message): pass
@@ -1516,7 +1550,7 @@ class NodeConnCB(object):
     def on_getheaders(self, conn, message): pass
     def on_headers(self, conn, message): pass
     def on_mempool(self, conn): pass
-    def on_open(self, conn): pass
+    def on_pong(self, conn, message): pass
     def on_reject(self, conn, message): pass
     def on_sendcmpct(self, conn, message): pass
     def on_sendheaders(self, conn, message): pass
@@ -1534,9 +1568,6 @@ class NodeConnCB(object):
         if conn.ver_send > BIP0031_VERSION:
             conn.send_message(msg_pong(message.nonce))
 
-    def on_pong(self, conn, message):
-        self.last_pong = message
-
     def on_verack(self, conn, message):
         conn.ver_recv = conn.ver_send
         self.verack_received = True
@@ -1549,15 +1580,44 @@ class NodeConnCB(object):
             conn.ver_recv = conn.ver_send
         conn.nServices = message.nServices
 
-    # Helper functions
-    ##################
+    # Connection helper methods
 
     def add_connection(self, conn):
         self.connection = conn
 
-    # Wrapper for the NodeConn's send_message function
+    def wait_for_disconnect(self, timeout=60):
+        test_function = lambda: not self.connected
+        assert wait_until(test_function, timeout=timeout)
+
+    # Message receiving helper methods
+
+    def wait_for_block(self, blockhash, timeout=60):
+        test_function = lambda: self.last_message.get("block") and self.last_message["block"].block.rehash() == blockhash
+        assert wait_until(test_function, timeout=timeout)
+
+    def wait_for_getdata(self, timeout=60):
+        test_function = lambda: self.last_message.get("getdata")
+        assert wait_until(test_function, timeout=timeout)
+
+    def wait_for_getheaders(self, timeout=60):
+        test_function = lambda: self.last_message.get("getheaders")
+        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
+        assert wait_until(test_function, timeout=timeout)
+
+    def wait_for_verack(self, timeout=60):
+        test_function = lambda: self.message_count["verack"]
+        assert wait_until(test_function, timeout=timeout)
+
+    # Message sending helper functions
+
     def send_message(self, message):
-        self.connection.send_message(message)
+        if self.connection:
+            self.connection.send_message(message)
+        else:
+            logger.error("Cannot send message. No connection to node!")
 
     def send_and_ping(self, message):
         self.send_message(message)
@@ -1565,27 +1625,11 @@ class NodeConnCB(object):
 
     # Sync up with the node
     def sync_with_ping(self, timeout=60):
-        def received_pong():
-            return (self.last_pong.nonce == self.ping_counter)
         self.send_message(msg_ping(nonce=self.ping_counter))
-        success = wait_until(received_pong, timeout=timeout)
-        if not success:
-            logger.error("sync_with_ping failed!")
-            raise AssertionError("sync_with_ping failed!")
+        test_function = lambda: self.last_message.get("pong") and self.last_message["pong"].nonce == self.ping_counter
+        assert wait_until(test_function, timeout=timeout)
         self.ping_counter += 1
-
-        return success
-
-    # Spin until verack message is received from the node.
-    # Tests may want to use this as a signal that the test can begin.
-    # This can be called from the testing thread, so it needs to acquire the
-    # global lock.
-    def wait_for_verack(self):
-        while True:
-            with mininode_lock:
-                if self.verack_received:
-                    return
-            time.sleep(0.05)
+        return True
 
 # The actual NodeConn class
 # This class provides an interface for a p2p connection to a specified node