p2p: Make timeout mockable and type safe, speed up test

1 files changed, 16 insertions, 14 deletions

diff --git a/test/functional/p2p_timeouts.py b/test/functional/p2p_timeouts.py
index a7e240dcfa..cf714bc888 100755
--- a/test/functional/p2p_timeouts.py
+++ b/test/functional/p2p_timeouts.py
@@ -4,13 +4,12 @@
 # file COPYING or http://www.opensource.org/licenses/mit-license.php.
 """Test various net timeouts.
 
-- Create three bitcoind nodes:
+- Create three peers:
 
     no_verack_node - we never send a verack in response to their version
     no_version_node - we never send a version (only a ping)
     no_send_node - we never send any P2P message.
 
-- Start all three nodes
 - Wait 1 second
 - Assert that we're connected
 - Send a ping to no_verack_node and no_version_node
@@ -21,12 +20,12 @@
 - Assert that we're no longer connected (timeout to receive version/verack is 3 seconds)
 """
 
-from time import sleep
-
 from test_framework.messages import msg_ping
 from test_framework.p2p import P2PInterface
 from test_framework.test_framework import BitcoinTestFramework
 
+import time
+
 
 class TestP2PConn(P2PInterface):
     def on_version(self, message):
@@ -41,7 +40,14 @@ class TimeoutsTest(BitcoinTestFramework):
         # set timeout to receive version/verack to 3 seconds
         self.extra_args = [["-peertimeout=3"]]
 
+    def mock_forward(self, delta):
+        self.mock_time += delta
+        self.nodes[0].setmocktime(self.mock_time)
+
     def run_test(self):
+        self.mock_time = int(time.time())
+        self.mock_forward(0)
+
         # Setup the p2p connections
         no_verack_node = self.nodes[0].add_p2p_connection(TestP2PConn(), wait_for_verack=False)
         no_version_node = self.nodes[0].add_p2p_connection(TestP2PConn(), send_version=False, wait_for_verack=False)
@@ -51,7 +57,7 @@ class TimeoutsTest(BitcoinTestFramework):
         # verack, since we never sent one
         no_verack_node.wait_for_verack()
 
-        sleep(1)
+        self.mock_forward(1)
 
         assert no_verack_node.is_connected
         assert no_version_node.is_connected
@@ -62,7 +68,7 @@ class TimeoutsTest(BitcoinTestFramework):
         with self.nodes[0].assert_debug_log(['non-version message before version handshake. Message "ping" from peer=1']):
             no_version_node.send_message(msg_ping())
 
-        sleep(1)
+        self.mock_forward(1)
 
         assert "version" in no_verack_node.last_message
 
@@ -80,14 +86,10 @@ class TimeoutsTest(BitcoinTestFramework):
         ]
 
         with self.nodes[0].assert_debug_log(expected_msgs=expected_timeout_logs):
-            sleep(3)
-            # By now, we waited a total of 5 seconds. Off-by-two for two
-            # reasons:
-            #  * The internal precision is one second
-            #  * Account for network delay
-            assert not no_verack_node.is_connected
-            assert not no_version_node.is_connected
-            assert not no_send_node.is_connected
+            self.mock_forward(2)
+            no_verack_node.wait_for_disconnect(timeout=1)
+            no_version_node.wait_for_disconnect(timeout=1)
+            no_send_node.wait_for_disconnect(timeout=1)
 
 
 if __name__ == '__main__':