test: Fix intermittent issues in feature_assumevalid.py

diff --git a/test/functional/feature_assumevalid.py b/test/functional/feature_assumevalid.py
index 2ee23c6e..4eaa74b3 100755
--- a/test/functional/feature_assumevalid.py
+++ b/test/functional/feature_assumevalid.py
@@ -79,17 +79,6 @@ class AssumeValidTest(BitcoinTestFramework):
         # signature so we can pass in the block hash as assumevalid.
         self.start_node(0)
 
-    def send_blocks_until_disconnected(self, p2p_conn):
-        """Keep sending blocks to the node until we're disconnected."""
-        for i in range(len(self.blocks)):
-            if not p2p_conn.is_connected:
-                break
-            try:
-                p2p_conn.send_without_ping(msg_block(self.blocks[i]))
-            except IOError:
-                assert not p2p_conn.is_connected
-                break
-
     def run_test(self):
         # Build the blockchain
         self.tip = int(self.nodes[0].getbestblockhash(), 16)
@@ -149,96 +138,95 @@ class AssumeValidTest(BitcoinTestFramework):
         self.start_node(4, extra_args=[f"-assumevalid={block102.hash_hex}"])
         self.start_node(5)
 
-
         # nodes[0]
-        # Send blocks to node0. Block 102 will be rejected.
+        self.log.info("Send blocks to node0. Block 102 will be rejected.")
+        p2p0 = self.nodes[0].add_p2p_connection(BaseNode())
+        p2p0.send_header_for_blocks(self.blocks[0:2000])
+        p2p0.send_header_for_blocks(self.blocks[2000:])
         with self.nodes[0].assert_debug_log(expected_msgs=[
             f"Enabling script verification at block #1 ({block_1_hash}): assumevalid=0 (always verify).",
+        ]):
+            p2p0.send_and_ping(msg_block(self.blocks[0]))
+        with self.nodes[0].assert_debug_log(expected_msgs=[
             "Block validation error: block-script-verify-flag-failed",
         ]):
-            p2p0 = self.nodes[0].add_p2p_connection(BaseNode())
-
-            p2p0.send_header_for_blocks(self.blocks[0:2000])
-            p2p0.send_header_for_blocks(self.blocks[2000:])
-
-            self.send_blocks_until_disconnected(p2p0)
-            self.wait_until(lambda: self.nodes[0].getblockcount() >= COINBASE_MATURITY + 1)
+            for i in range(1, 103):
+                p2p0.send_without_ping(msg_block(self.blocks[i]))
+            p2p0.wait_for_disconnect()
             assert_equal(self.nodes[0].getblockcount(), COINBASE_MATURITY + 1)
-            self.wait_until(lambda: next(filter(lambda x: x["hash"] == self.blocks[-1].hash_hex, self.nodes[0].getchaintips()))["status"] == "invalid")
-
+            assert_equal(next(filter(lambda x: x["hash"] == self.blocks[-1].hash_hex, self.nodes[0].getchaintips()))["status"], "invalid")
 
         # nodes[1]
+        self.log.info("Send all blocks to node1. All blocks will be accepted.")
+        p2p1 = self.nodes[1].add_p2p_connection(BaseNode())
+        p2p1.send_header_for_blocks(self.blocks[0:2000])
+        p2p1.send_header_for_blocks(self.blocks[2000:])
         with self.nodes[1].assert_debug_log(expected_msgs=[
             f"Disabling script verification at block #1 ({self.blocks[0].hash_hex}).",
+        ]):
+            p2p1.send_and_ping(msg_block(self.blocks[0]))
+        with self.nodes[1].assert_debug_log(expected_msgs=[
             f"Enabling script verification at block #103 ({self.blocks[102].hash_hex}): block height above assumevalid height.",
         ]):
-            p2p1 = self.nodes[1].add_p2p_connection(BaseNode())
-
-            p2p1.send_header_for_blocks(self.blocks[0:2000])
-            p2p1.send_header_for_blocks(self.blocks[2000:])
-            # Send all blocks to node1. All blocks will be accepted.
-            for i in range(2202):
+            for i in range(1, 2202):
                 p2p1.send_without_ping(msg_block(self.blocks[i]))
             # Syncing 2200 blocks can take a while on slow systems. Give it plenty of time to sync.
             p2p1.sync_with_ping(timeout=960)
-            assert_equal(self.nodes[1].getblock(self.nodes[1].getbestblockhash())['height'], 2202)
-
+            assert_equal(self.nodes[1].getblockcount(), 2202)
 
         # nodes[2]
-        # Send blocks to node2. Block 102 will be rejected.
+        self.log.info("Send blocks to node2. Block 102 will be rejected.")
+        p2p2 = self.nodes[2].add_p2p_connection(BaseNode())
+        p2p2.send_header_for_blocks(self.blocks[0:200])
         with self.nodes[2].assert_debug_log(expected_msgs=[
             f"Enabling script verification at block #1 ({block_1_hash}): block too recent relative to best header.",
+        ]):
+            p2p2.send_and_ping(msg_block(self.blocks[0]))
+        with self.nodes[2].assert_debug_log(expected_msgs=[
             "Block validation error: block-script-verify-flag-failed",
         ]):
-            p2p2 = self.nodes[2].add_p2p_connection(BaseNode())
-            p2p2.send_header_for_blocks(self.blocks[0:200])
-
-            self.send_blocks_until_disconnected(p2p2)
-
-            self.wait_until(lambda: self.nodes[2].getblockcount() >= COINBASE_MATURITY + 1)
+            for i in range(1, 103):
+                p2p2.send_without_ping(msg_block(self.blocks[i]))
+            p2p2.wait_for_disconnect()
             assert_equal(self.nodes[2].getblockcount(), COINBASE_MATURITY + 1)
-
+            assert_equal(next(filter(lambda x: x["hash"] == self.blocks[199].hash_hex, self.nodes[2].getchaintips()))["status"], "invalid")
 
         # nodes[3]
+        self.log.info("Send two header chains, and a block not in the best header chain to node3.")
+        best_hash = self.nodes[3].getbestblockhash()
+        tip_block = self.nodes[3].getblock(best_hash)
+        second_chain_tip, second_chain_time, second_chain_height = int(best_hash, 16), tip_block["time"] + 1, tip_block["height"] + 1
+        second_chain = []
+        for _ in range(150):
+            block = create_block(second_chain_tip, create_coinbase(second_chain_height), second_chain_time)
+            block.solve()
+            second_chain.append(block)
+            second_chain_tip, second_chain_time, second_chain_height = block.hash_int, second_chain_time + 1, second_chain_height + 1
+        p2p3 = self.nodes[3].add_p2p_connection(BaseNode())
+        p2p3.send_header_for_blocks(second_chain)
+        p2p3.send_header_for_blocks(self.blocks[0:103])
         with self.nodes[3].assert_debug_log(expected_msgs=[
             f"Enabling script verification at block #1 ({block_1_hash}): block not in best header chain.",
         ]):
-            best_hash = self.nodes[3].getbestblockhash()
-            tip_block = self.nodes[3].getblock(best_hash)
-            second_chain_tip, second_chain_time, second_chain_height = int(best_hash, 16), tip_block["time"] + 1, tip_block["height"] + 1
-            second_chain = []
-            for _ in range(150):
-                block = create_block(second_chain_tip, create_coinbase(second_chain_height), second_chain_time)
-                block.solve()
-                second_chain.append(block)
-                second_chain_tip, second_chain_time, second_chain_height = block.hash_int, second_chain_time + 1, second_chain_height + 1
-
-            p2p3 = self.nodes[3].add_p2p_connection(BaseNode())
-
-            p2p3.send_header_for_blocks(second_chain)
-            p2p3.send_header_for_blocks(self.blocks[0:103])
-
-            p2p3.send_without_ping(msg_block(self.blocks[0]))
-            self.wait_until(lambda: self.nodes[3].getblockcount() == 1)
-
+            p2p3.send_and_ping(msg_block(self.blocks[0]))
+            assert_equal(self.nodes[3].getblockcount(), 1)
 
         # nodes[4]
+        self.log.info("Send a block not in the assumevalid header chain to node4.")
         genesis_hash = self.nodes[4].getbestblockhash()
         genesis_time = self.nodes[4].getblock(genesis_hash)['time']
         alt1 = create_block(int(genesis_hash, 16), create_coinbase(1), genesis_time + 2)
         alt1.solve()
+        p2p4 = self.nodes[4].add_p2p_connection(BaseNode())
+        p2p4.send_header_for_blocks(self.blocks[0:103])
         with self.nodes[4].assert_debug_log(expected_msgs=[
             f"Enabling script verification at block #1 ({alt1.hash_hex}): block not in assumevalid chain.",
         ]):
-            p2p4 = self.nodes[4].add_p2p_connection(BaseNode())
-            p2p4.send_header_for_blocks(self.blocks[0:103])
-
-            p2p4.send_without_ping(msg_block(alt1))
-            self.wait_until(lambda: self.nodes[4].getblockcount() == 1)
-
+            p2p4.send_and_ping(msg_block(alt1))
+            assert_equal(self.nodes[4].getblockcount(), 1)
 
         # nodes[5]
-        # Reindex to hit specific assumevalid gates (no races with header downloads/chainwork during startup).
+        self.log.info("Reindex to hit specific assumevalid gates (no races with header downloads/chainwork during startup).")
         p2p5 = self.nodes[5].add_p2p_connection(BaseNode())
         p2p5.send_header_for_blocks(self.blocks[0:200])
         p2p5.send_without_ping(msg_block(self.blocks[0]))