Repository: Chia-Network/chia-blockchain
File: chia/_tests/wallet/sync/test_wallet_sync.py (Python)

Code before:
    await add_blocks_in_batches(default_400_blocks, trusted_full_node_api.full_node)

    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),
        caplog.at_level(logging.INFO),
    ):
        # Connect to the untrusted peer and wait until the long sync started
        assert await wallet_server.start_client(PeerInfo(self_hostname, untrusted_full_node_server.get_port()), None)
        untrusted_conn = wallet_server.all_connections.get(untrusted_full_node_server.node_id)
        assert untrusted_conn is not None
        await time_out_assert(30, wallet_syncing)

        # 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
        assert (
            f"Connection closed: {untrusted_conn.peer_info.host}, node id: {untrusted_full_node_server.node_id.hex()}"
            in caplog.text
        )
        # And that we only have a trusted peer left
        await time_out_assert(30, only_trusted_peer)


@pytest.mark.anyio
@pytest.mark.parametrize("chain_length", [0, 100])
@pytest.mark.parametrize("fork_point", [500, 1500])
# TODO: todo_v2_plots once we have new test chains, we can probably re-enable this
@pytest.mark.limit_consensus_modes(
    allows=[ConsensusMode.PLAIN, ConsensusMode.HARD_FORK_2_0],
    reason="after plot-v1 phase-out, the chains aren't valid anymore",
)
async def test_long_reorg_nodes_and_wallet(
    chain_length: int,
    fork_point: int,
    three_nodes: list[FullNodeAPI],
    simulator_and_wallet: OldSimulatorsAndWallets,
    default_10000_blocks: list[FullBlock],
    test_long_reorg_blocks: list[FullBlock],
    test_long_reorg_1500_blocks: list[FullBlock],
    self_hostname: str,
) -> None:
    full_node_1, full_node_2, _ = three_nodes
    _, [wallet], _ = simulator_and_wallet
    wallet_node = wallet[0]
    wallet_server = wallet[1]

Review comment (suggestion):
From what I see I think they just had a thinko about asserting `only_trusted_peer`'s return (which happens by default in `time_out_assert`) and asserting `time_out_assert` instead.

Code after:
    await add_blocks_in_batches(default_1500_blocks[:600], trusted_full_node_api.full_node)
    await add_blocks_in_batches(default_1000_blocks[:600], untrusted_full_node_api.full_node)

    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
        assert await wallet_server.start_client(PeerInfo(self_hostname, untrusted_full_node_server.get_port()), None)
        untrusted_conn = wallet_server.all_connections.get(untrusted_full_node_server.node_id)
        assert untrusted_conn is not None
        await time_out_assert(30, wallet_syncing)

        # 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, 599)
        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
        assert (
            f"Connection closed: {untrusted_conn.peer_info.host}, node id: {untrusted_full_node_server.node_id}"
            in caplog.text
        )
        # And that we only have a trusted peer left
        await time_out_assert(30, only_trusted_peer)


@pytest.mark.anyio
@pytest.mark.parametrize("chain_length", [0, 100])
@pytest.mark.parametrize("fork_point", [500, 1500])
# TODO: todo_v2_plots once we have new test chains, we can probably re-enable this
@pytest.mark.limit_consensus_modes(
    allows=[ConsensusMode.PLAIN, ConsensusMode.HARD_FORK_2_0],
    reason="after plot-v1 phase-out, the chains aren't valid anymore",
)
async def test_long_reorg_nodes_and_wallet(
    chain_length: int,
    fork_point: int,
    three_nodes: list[FullNodeAPI],
    simulator_and_wallet: OldSimulatorsAndWallets,
    default_10000_blocks: list[FullBlock],
    test_long_reorg_blocks: list[FullBlock],
    test_long_reorg_1500_blocks: list[FullBlock],
    self_hostname: str,
) -> None:
    full_node_1, full_node_2, _ = three_nodes
    _, [wallet], _ = simulator_and_wallet
    wallet_node = wallet[0]
    wallet_server = wallet[1]
    # Trusted node sync