[tests] Add p2p connection to TestNode

p2p connections can now be added to TestNode instances.

This commit also updates the example test to use the new
p2p interface in TestNode to demonstrate usage.

A future commit will update the existing tests to use p2p through the
TestNode.

2 files changed, 50 insertions, 18 deletions

diff --git a/test/functional/example_test.py b/test/functional/example_test.py
index 87d73ad14a..ba40f33016 100755
--- a/test/functional/example_test.py
+++ b/test/functional/example_test.py
@@ -18,7 +18,6 @@ from test_framework.blocktools import (create_block, create_coinbase)
 from test_framework.mininode import (
     CInv,
     NetworkThread,
-    NodeConn,
     NodeConnCB,
     mininode_lock,
     msg_block,
@@ -28,7 +27,6 @@ from test_framework.test_framework import BitcoinTestFramework
 from test_framework.util import (
     assert_equal,
     connect_nodes,
-    p2p_port,
     wait_until,
 )
 
@@ -134,16 +132,13 @@ class ExampleTest(BitcoinTestFramework):
         """Main test logic"""
 
         # Create a P2P connection to one of the nodes
-        node0 = BaseNode()
-        connections = []
-        connections.append(NodeConn('127.0.0.1', p2p_port(0), self.nodes[0], node0))
-        node0.add_connection(connections[0])
+        self.nodes[0].add_p2p_connection(BaseNode())
 
         # Start up network handling in another thread. This needs to be called
         # after the P2P connections have been created.
         NetworkThread().start()
         # wait_for_verack ensures that the P2P connection is fully up.
-        node0.wait_for_verack()
+        self.nodes[0].p2p.wait_for_verack()
 
         # Generating a block on one of the nodes will get us out of IBD
         blocks = [int(self.nodes[0].generate(nblocks=1)[0], 16)]
@@ -180,7 +175,7 @@ class ExampleTest(BitcoinTestFramework):
             block.solve()
             block_message = msg_block(block)
             # Send message is used to send a P2P message to the node over our NodeConn connection
-            node0.send_message(block_message)
+            self.nodes[0].p2p.send_message(block_message)
             self.tip = block.sha256
             blocks.append(self.tip)
             self.block_time += 1
@@ -193,28 +188,26 @@ class ExampleTest(BitcoinTestFramework):
         connect_nodes(self.nodes[1], 2)
 
         self.log.info("Add P2P connection to node2")
-        node2 = BaseNode()
-        connections.append(NodeConn('127.0.0.1', p2p_port(2), self.nodes[2], node2))
-        node2.add_connection(connections[1])
-        node2.wait_for_verack()
+        self.nodes[2].add_p2p_connection(BaseNode())
+        self.nodes[2].p2p.wait_for_verack()
 
         self.log.info("Wait for node2 reach current tip. Test that it has propagated all the blocks to us")
 
         getdata_request = msg_getdata()
         for block in blocks:
             getdata_request.inv.append(CInv(2, block))
-        node2.send_message(getdata_request)
+        self.nodes[2].p2p.send_message(getdata_request)
 
         # wait_until() will loop until a predicate condition is met. Use it to test properties of the
         # NodeConnCB objects.
-        wait_until(lambda: sorted(blocks) == sorted(list(node2.block_receive_map.keys())), timeout=5, lock=mininode_lock)
+        wait_until(lambda: sorted(blocks) == sorted(list(self.nodes[2].p2p.block_receive_map.keys())), timeout=5, lock=mininode_lock)
 
         self.log.info("Check that each block was received only once")
         # The network thread uses a global lock on data access to the NodeConn objects when sending and receiving
         # messages. The test thread should acquire the global lock before accessing any NodeConn data to avoid locking
         # and synchronization issues. Note wait_until() acquires this global lock when testing the predicate.
         with mininode_lock:
-            for block in node2.block_receive_map.values():
+            for block in self.nodes[2].p2p.block_receive_map.values():
                 assert_equal(block, 1)
 
 if __name__ == '__main__':
diff --git a/test/functional/test_framework/test_node.py b/test/functional/test_framework/test_node.py
index 41c31c2d3d..8b28064c46 100755
--- a/test/functional/test_framework/test_node.py
+++ b/test/functional/test_framework/test_node.py
@@ -13,13 +13,15 @@ import os
 import subprocess
 import time
 
+from .authproxy import JSONRPCException
+from .mininode import NodeConn
 from .util import (
     assert_equal,
     get_rpc_proxy,
     rpc_url,
     wait_until,
+    p2p_port,
 )
-from .authproxy import JSONRPCException
 
 BITCOIND_PROC_WAIT_TIMEOUT = 60
 
@@ -31,9 +33,11 @@ class TestNode():
     - state about the node (whether it's running, etc)
     - a Python subprocess.Popen object representing the running process
     - an RPC connection to the node
+    - one or more P2P connections to the node
+
 
-    To make things easier for the test writer, a bit of magic is happening under the covers.
-    Any unrecognised messages will be dispatched to the RPC connection."""
+    To make things easier for the test writer, any unrecognised messages will
+    be dispatched to the RPC connection."""
 
     def __init__(self, i, dirname, extra_args, rpchost, timewait, binary, stderr, mocktime, coverage_dir):
         self.index = i
@@ -63,6 +67,8 @@ class TestNode():
         self.url = None
         self.log = logging.getLogger('TestFramework.node%d' % i)
 
+        self.p2ps = []
+
     def __getattr__(self, name):
         """Dispatches any unrecognised messages to the RPC connection."""
         assert self.rpc_connected and self.rpc is not None, "Error: no RPC connection"
@@ -119,6 +125,7 @@ class TestNode():
             self.stop()
         except http.client.CannotSendRequest:
             self.log.exception("Unable to stop node.")
+        del self.p2ps[:]
 
     def is_node_stopped(self):
         """Checks whether the node has stopped.
@@ -151,6 +158,38 @@ class TestNode():
         self.encryptwallet(passphrase)
         self.wait_until_stopped()
 
+    def add_p2p_connection(self, p2p_conn, **kwargs):
+        """Add a p2p connection to the node.
+
+        This method adds the p2p connection to the self.p2ps list and also
+        returns the connection to the caller."""
+        if 'dstport' not in kwargs:
+            kwargs['dstport'] = p2p_port(self.index)
+        if 'dstaddr' not in kwargs:
+            kwargs['dstaddr'] = '127.0.0.1'
+        self.p2ps.append(p2p_conn)
+        kwargs.update({'rpc': self.rpc, 'callback': p2p_conn})
+        p2p_conn.add_connection(NodeConn(**kwargs))
+
+        return p2p_conn
+
+    @property
+    def p2p(self):
+        """Return the first p2p connection
+
+        Convenience property - most tests only use a single p2p connection to each
+        node, so this saves having to write node.p2ps[0] many times."""
+        assert self.p2ps, "No p2p connection"
+        return self.p2ps[0]
+
+    def disconnect_p2p(self, index=0):
+        """Close the p2p connection to the node."""
+        # Connection could have already been closed by other end. Calling disconnect_p2p()
+        # on an already disconnected p2p connection is not an error.
+        if self.p2ps[index].connection is not None:
+            self.p2ps[index].connection.disconnect_node()
+        del self.p2ps[index]
+
 class TestNodeCLI():
     """Interface to bitcoin-cli for an individual node"""