test: make p2p_ibd_stalling.py more modular

This is in preparation to add more subtests.
Also adjust waiting condition to replace the
total_bytes_recv_for_blocks magical number hack.

test/functional/p2p_ibd_stalling.py

@@ -49,41 +49,42 @@ class P2PIBDStallingTest(BitcoinTestFramework):
         self.setup_clean_chain = True
         self.num_nodes = 1
 
-    def run_test(self):
-        NUM_BLOCKS = 1025
-        NUM_PEERS = 4
+    def prepare_blocks(self):
+        self.log.info("Prepare blocks without sending them to any node")
+        self.NUM_BLOCKS = 1025
+        self.block_dict = {}
+        self.blocks = []
+
         node = self.nodes[0]
         tip = int(node.getbestblockhash(), 16)
-        blocks = []
         height = 1
         block_time = node.getblock(node.getbestblockhash())['time'] + 1
-        self.log.info("Prepare blocks without sending them to the node")
-        block_dict = {}
-        for _ in range(NUM_BLOCKS):
-            blocks.append(create_block(tip, create_coinbase(height), block_time))
-            blocks[-1].solve()
-            tip = blocks[-1].sha256
+        for _ in range(self.NUM_BLOCKS):
+            self.blocks.append(create_block(tip, create_coinbase(height), block_time))
+            self.blocks[-1].solve()
+            tip = self.blocks[-1].sha256
             block_time += 1
             height += 1
-            block_dict[blocks[-1].sha256] = blocks[-1]
-        stall_block = blocks[0].sha256
+            self.block_dict[self.blocks[-1].sha256] = self.blocks[-1]
+
+    def ibd_stalling(self):
+        NUM_PEERS = 4
+        stall_block = self.blocks[0].sha256
+        node = self.nodes[0]
 
         headers_message = msg_headers()
-        headers_message.headers = [CBlockHeader(b) for b in blocks[:NUM_BLOCKS-1]]
+        headers_message.headers = [CBlockHeader(b) for b in self.blocks[:self.NUM_BLOCKS-1]]
         peers = []
 
         self.log.info("Check that a staller does not get disconnected if the 1024 block lookahead buffer is filled")
         self.mocktime = int(time.time()) + 1
         for id in range(NUM_PEERS):
             peers.append(node.add_outbound_p2p_connection(P2PStaller(stall_block), p2p_idx=id, connection_type="outbound-full-relay"))
-            peers[-1].block_store = block_dict
+            peers[-1].block_store = self.block_dict
             peers[-1].send_message(headers_message)
 
-        # Need to wait until 1023 blocks are received - the magic total bytes number is a workaround in lack of an rpc
-        # returning the number of downloaded (but not connected) blocks.
-        bytes_recv = 172761 if not self.options.v2transport else 169692
-        self.wait_until(lambda: self.total_bytes_recv_for_blocks() == bytes_recv)
-
+        # Wait until all blocks are received (except for stall_block), so that no other blocks are in flight.
+        self.wait_until(lambda: sum(len(peer['inflight']) for peer in node.getpeerinfo()) == 1)
         self.all_sync_send_with_ping(peers)
         # If there was a peer marked for stalling, it would get disconnected
         self.mocktime += 3
@@ -92,7 +93,7 @@ class P2PIBDStallingTest(BitcoinTestFramework):
         assert_equal(node.num_test_p2p_connections(), NUM_PEERS)
 
         self.log.info("Check that increasing the window beyond 1024 blocks triggers stalling logic")
-        headers_message.headers = [CBlockHeader(b) for b in blocks]
+        headers_message.headers = [CBlockHeader(b) for b in self.blocks]
         with node.assert_debug_log(expected_msgs=['Stall started']):
             for p in peers:
                 p.send_message(headers_message)
@@ -138,17 +139,10 @@ class P2PIBDStallingTest(BitcoinTestFramework):
         with node.assert_debug_log(expected_msgs=['Decreased stalling timeout to 2 seconds']):
             for p in peers:
                 if p.is_connected and (stall_block in p.getdata_requests):
-                    p.send_message(msg_block(block_dict[stall_block]))
+                    p.send_message(msg_block(self.block_dict[stall_block]))
 
         self.log.info("Check that all outstanding blocks get connected")
-        self.wait_until(lambda: node.getblockcount() == NUM_BLOCKS)
-
-    def total_bytes_recv_for_blocks(self):
-        total = 0
-        for info in self.nodes[0].getpeerinfo():
-            if ("block" in info["bytesrecv_per_msg"].keys()):
-                total += info["bytesrecv_per_msg"]["block"]
-        return total
+        self.wait_until(lambda: node.getblockcount() == self.NUM_BLOCKS)
 
     def all_sync_send_with_ping(self, peers):
         for p in peers:
@@ -161,6 +155,10 @@ class P2PIBDStallingTest(BitcoinTestFramework):
                 return True
         return False
 
+    def run_test(self):
+        self.prepare_blocks()
+        self.ibd_stalling()
+
 
 if __name__ == '__main__':
     P2PIBDStallingTest(__file__).main()