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

1 files changed, 20 insertions, 79 deletions

diff --git a/test/functional/p2p-segwit.py b/test/functional/p2p-segwit.py
index 3dd78c8b7d..1bd7c90956 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