[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.

1 files changed, 42 insertions, 3 deletions

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"""