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