From c9919433994924931757aa7ec278c50744ede35d Mon Sep 17 00:00:00 2001
From: Amiti Uttarwar <amiti@uttarwar.org>
Date: Mon, 29 Mar 2021 15:19:37 -0700
Subject: [test] Refactor the addr relay test to prepare for new tests

Moves setting up the addr message into a repeatable function, and breaks up the
existing tests into separate functions for legibility.
---
 test/functional/p2p_addr_relay.py | 57 ++++++++++++++++++++++++++-------------
 1 file changed, 38 insertions(+), 19 deletions(-)

(limited to 'test/functional/p2p_addr_relay.py')

diff --git a/test/functional/p2p_addr_relay.py b/test/functional/p2p_addr_relay.py
index 69821763bd..b728ed428b 100755
--- a/test/functional/p2p_addr_relay.py
+++ b/test/functional/p2p_addr_relay.py
@@ -19,18 +19,6 @@ from test_framework.util import (
 )
 import time
 
-# Keep this with length <= 10. Addresses from larger messages are not relayed.
-ADDRS = []
-num_ipv4_addrs = 10
-
-for i in range(num_ipv4_addrs):
-    addr = CAddress()
-    addr.time = int(time.time()) + i
-    addr.nServices = NODE_NETWORK | NODE_WITNESS
-    addr.ip = "123.123.123.{}".format(i % 256)
-    addr.port = 8333 + i
-    ADDRS.append(addr)
-
 
 class AddrReceiver(P2PInterface):
     num_ipv4_received = 0
@@ -45,33 +33,62 @@ class AddrReceiver(P2PInterface):
 
 
 class AddrTest(BitcoinTestFramework):
+    counter = 0
+    mocktime = int(time.time())
+
     def set_test_params(self):
         self.num_nodes = 1
 
     def run_test(self):
-        self.log.info('Create connection that sends addr messages')
-        addr_source = self.nodes[0].add_p2p_connection(P2PInterface())
+        self.oversized_addr_test()
+        self.relay_tests()
+
+    def setup_addr_msg(self, num):
+        addrs = []
+        for i in range(num):
+            addr = CAddress()
+            addr.time = self.mocktime + i
+            addr.nServices = NODE_NETWORK | NODE_WITNESS
+            addr.ip = f"123.123.123.{self.counter % 256}"
+            addr.port = 8333 + i
+            addrs.append(addr)
+            self.counter += 1
+
         msg = msg_addr()
+        msg.addrs = addrs
+        return msg
 
-        self.log.info('Send too-large addr message')
-        msg.addrs = ADDRS * 101  # more than 1000 addresses in one message
+    def oversized_addr_test(self):
+        self.log.info('Send an addr message that is too large')
+        addr_source = self.nodes[0].add_p2p_connection(P2PInterface())
+
+        msg = self.setup_addr_msg(1010)
         with self.nodes[0].assert_debug_log(['addr message size = 1010']):
             addr_source.send_and_ping(msg)
 
+        self.nodes[0].disconnect_p2ps()
+
+    def relay_tests(self):
         self.log.info('Check that addr message content is relayed and added to addrman')
+        addr_source = self.nodes[0].add_p2p_connection(P2PInterface())
         num_receivers = 7
         receivers = []
         for _ in range(num_receivers):
             receivers.append(self.nodes[0].add_p2p_connection(AddrReceiver()))
-        msg.addrs = ADDRS
+
+        # Keep this with length <= 10. Addresses from larger messages are not
+        # relayed.
+        num_ipv4_addrs = 10
+        msg = self.setup_addr_msg(num_ipv4_addrs)
         with self.nodes[0].assert_debug_log(
             [
                 'Added {} addresses from 127.0.0.1: 0 tried'.format(num_ipv4_addrs),
-                'received: addr (301 bytes) peer=0',
+                'received: addr (301 bytes) peer=1',
             ]
         ):
             addr_source.send_and_ping(msg)
-            self.nodes[0].setmocktime(int(time.time()) + 30 * 60)
+            self.mocktime += 30 * 60
+            self.nodes[0].setmocktime(self.mocktime)
             for receiver in receivers:
                 receiver.sync_with_ping()
 
@@ -82,6 +99,8 @@ class AddrTest(BitcoinTestFramework):
         ipv4_branching_factor = 2
         assert_equal(total_ipv4_received, num_ipv4_addrs * ipv4_branching_factor)
 
+        self.nodes[0].disconnect_p2ps()
+
 
 if __name__ == '__main__':
     AddrTest().main()
-- 
cgit v1.2.3