Fixes to test_long_sync_untrusted_break

chia/_tests/wallet/sync/test_wallet_sync.py

@@ -1650,27 +1650,16 @@ async def test_long_sync_untrusted_break(
     wallet_node.config["trusted_peers"] = {trusted_full_node_server.node_id.hex(): None}
     wallet_node.config["use_delta_sync"] = use_delta_sync
 
-    sync_canceled = False
-
     async def register_for_ph_updates(
         self: object,
         request: wallet_protocol.RegisterForPhUpdates,
         peer: WSChiaConnection,
     ) -> None:
-        nonlocal sync_canceled
-        # Just sleep a long time here to simulate a long-running untrusted sync
-        try:
-            await asyncio.sleep(120)
-        except Exception:
-            sync_canceled = True
-            raise
+        await peer.wait_until_closed()
 
     def wallet_syncing() -> bool:
         return wallet_node.wallet_state_manager.sync_mode
 
-    def check_sync_canceled() -> bool:
-        return sync_canceled
-
     def synced_to_trusted() -> bool:
         return trusted_full_node_server.node_id in wallet_node.synced_peers
 
@@ -1683,22 +1672,26 @@ def only_trusted_peer() -> bool:
 
     await add_blocks_in_batches(default_1000_blocks[:400], untrusted_full_node_api.full_node)
 
-    with patch_request_handler(api=untrusted_full_node_api, handler=register_for_ph_updates):
+    with (
+        patch_request_handler(api=untrusted_full_node_api, handler=register_for_ph_updates),
+        caplog.at_level(logging.INFO),
+    ):
         # Connect to the untrusted peer and wait until the long sync started
-        await wallet_server.start_client(PeerInfo(self_hostname, untrusted_full_node_server.get_port()), None)
+        assert await wallet_server.start_client(PeerInfo(self_hostname, untrusted_full_node_server.get_port()), None)
         await time_out_assert(30, wallet_syncing)
-        with caplog.at_level(logging.INFO):
-            # Connect to the trusted peer and make sure the running untrusted long sync gets interrupted via disconnect
-            await wallet_server.start_client(PeerInfo(self_hostname, trusted_full_node_server.get_port()), None)
-            await time_out_assert(600, wallet_height_at_least, True, wallet_node, len(default_400_blocks) - 1)
-            assert time_out_assert(10, synced_to_trusted)
-            assert untrusted_full_node_server.node_id not in wallet_node.synced_peers
-            assert "Connected to a synced trusted peer, disconnecting from all untrusted nodes." in caplog.text
-
-        # Make sure the sync was interrupted
-        assert time_out_assert(30, check_sync_canceled)
+
+        # Connect to the trusted peer and make sure the running untrusted long sync gets interrupted via disconnect
+        assert await wallet_server.start_client(PeerInfo(self_hostname, trusted_full_node_server.get_port()), None)
+        await time_out_assert(600, wallet_height_at_least, True, wallet_node, len(default_400_blocks) - 1)
+        await time_out_assert(10, synced_to_trusted)
+        assert untrusted_full_node_server.node_id not in wallet_node.synced_peers
+
+        assert "Connected to a synced trusted peer, disconnecting from all untrusted nodes." in caplog.text
+        # check that we disconnected from the untrusted peer
+        assert f"Connection closed: 127.0.0.1, node id: {untrusted_full_node_server.node_id.hex()}" in caplog.text
+
         # And that we only have a trusted peer left
-        assert time_out_assert(30, only_trusted_peer)
+        await time_out_assert(30, only_trusted_peer)
 
 
 @pytest.mark.anyio