chia-blockchain/chia/_tests/wallet/test_wallet_state_manager.py

from __future__ import annotations

import logging
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from dataclasses import dataclass
from types import SimpleNamespace
from typing import cast

import pytest
from chia_rs import CoinState, G2Element
from chia_rs.sized_bytes import bytes32
from chia_rs.sized_ints import uint32, uint64

from chia._tests.conftest import ConsensusMode
from chia._tests.environments.wallet import WalletStateTransition, WalletTestFramework
from chia._tests.util.setup_nodes import OldSimulatorsAndWallets
from chia._tests.util.time_out_assert import time_out_assert
from chia.protocols.outbound_message import NodeType
from chia.types.blockchain_format.coin import Coin
from chia.types.blockchain_format.program import Program
from chia.types.coin_spend import make_spend
from chia.types.peer_info import PeerInfo
from chia.wallet import wallet_state_manager as wsm_mod
from chia.wallet.derivation_record import DerivationRecord
from chia.wallet.derive_keys import master_sk_to_wallet_sk, master_sk_to_wallet_sk_unhardened
from chia.wallet.nft_wallet.nft_wallet import NFTWallet
from chia.wallet.nft_wallet.uncurry_nft import NFTCoinData, UncurriedNFT
from chia.wallet.remote_wallet.remote_wallet import RemoteWallet
from chia.wallet.transaction_record import TransactionRecord
from chia.wallet.util.transaction_type import TransactionType
from chia.wallet.util.wallet_types import WalletType
from chia.wallet.wallet_request_types import (
    CreateNewWallet,
    CreateNewWalletType,
    ExtendDerivationIndex,
    GetCoinRecordsByNames,
    GetHeightInfo,
    GetHeightInfoResponse,
    GetPuzzleAndSolution,
    GetPuzzleAndSolutionResponse,
    GetSpendableCoins,
    GetWalletBalance,
    PushTransactions,
    SelectCoins,
)
from chia.wallet.wallet_rpc_api import MAX_DERIVATION_INDEX_DELTA
from chia.wallet.wallet_spend_bundle import WalletSpendBundle
from chia.wallet.wallet_state_manager import SyncStatus, WalletStateManager


@asynccontextmanager
async def assert_sync_mode(wallet_state_manager: WalletStateManager, target_height: uint32) -> AsyncIterator[None]:
    assert not wallet_state_manager.lock.locked()
    assert not wallet_state_manager.sync_mode
    assert wallet_state_manager.sync_target is None
    new_current_height = max(0, target_height - 1)
    await wallet_state_manager.blockchain.set_finished_sync_up_to(new_current_height)
    async with wallet_state_manager.set_sync_mode(target_height) as current_height:
        assert current_height == new_current_height
        assert wallet_state_manager.sync_mode
        assert wallet_state_manager.lock.locked()
        assert wallet_state_manager.sync_target == target_height
        yield
    assert not wallet_state_manager.lock.locked()
    assert not wallet_state_manager.sync_mode
    assert wallet_state_manager.sync_target is None


@pytest.mark.limit_consensus_modes(allowed=[ConsensusMode.HARD_FORK_2_0])
@pytest.mark.anyio
async def test_set_sync_mode(simulator_and_wallet: OldSimulatorsAndWallets) -> None:
    _, [(wallet_node, _)], _ = simulator_and_wallet
    async with assert_sync_mode(wallet_node.wallet_state_manager, uint32(1)):
        pass
    async with assert_sync_mode(wallet_node.wallet_state_manager, uint32(22)):
        pass
    async with assert_sync_mode(wallet_node.wallet_state_manager, uint32(333)):
        pass


@pytest.mark.limit_consensus_modes(allowed=[ConsensusMode.HARD_FORK_2_0])
@pytest.mark.anyio
async def test_set_sync_mode_exception(simulator_and_wallet: OldSimulatorsAndWallets) -> None:
    _, [(wallet_node, _)], _ = simulator_and_wallet
    async with assert_sync_mode(wallet_node.wallet_state_manager, uint32(1)):
        raise Exception


@pytest.mark.parametrize("hardened", [True, False])
@pytest.mark.limit_consensus_modes(allowed=[ConsensusMode.HARD_FORK_2_0])
@pytest.mark.anyio
async def test_get_private_key(simulator_and_wallet: OldSimulatorsAndWallets, hardened: bool) -> None:
    _, [(wallet_node, _)], _ = simulator_and_wallet
    wallet_state_manager: WalletStateManager = wallet_node.wallet_state_manager
    derivation_index = uint32(10000)
    conversion_method = master_sk_to_wallet_sk if hardened else master_sk_to_wallet_sk_unhardened
    expected_private_key = conversion_method(wallet_state_manager.get_master_private_key(), derivation_index)
    record = DerivationRecord(
        derivation_index,
        bytes32(b"0" * 32),
        expected_private_key.get_g1(),
        WalletType.STANDARD_WALLET,
        uint32(1),
        hardened,
    )
    await wallet_state_manager.puzzle_store.add_derivation_paths([record])
    assert await wallet_state_manager.get_private_key(record.puzzle_hash) == expected_private_key


@pytest.mark.limit_consensus_modes(allowed=[ConsensusMode.HARD_FORK_2_0])
@pytest.mark.anyio
async def test_get_private_key_failure(simulator_and_wallet: OldSimulatorsAndWallets) -> None:
    _, [(wallet_node, _)], _ = simulator_and_wallet
    wallet_state_manager: WalletStateManager = wallet_node.wallet_state_manager
    invalid_puzzle_hash = bytes32(b"1" * 32)
    with pytest.raises(ValueError, match=f"No key for puzzle hash: {invalid_puzzle_hash.hex()}"):
        await wallet_state_manager.get_private_key(bytes32(b"1" * 32))


@pytest.mark.limit_consensus_modes(allowed=[ConsensusMode.HARD_FORK_2_0])
@pytest.mark.anyio
async def test_determine_coin_type(simulator_and_wallet: OldSimulatorsAndWallets, self_hostname: str) -> None:
    full_nodes, wallets, _ = simulator_and_wallet
    full_node_api = full_nodes[0]
    full_node_server = full_node_api.full_node.server
    wallet_node, wallet_server = wallets[0]
    await wallet_server.start_client(PeerInfo(self_hostname, full_node_server.get_port()), None)
    wallet_state_manager: WalletStateManager = wallet_node.wallet_state_manager
    peer = wallet_node.server.get_connections(NodeType.FULL_NODE)[0]
    assert (None, None) == await wallet_state_manager.determine_coin_type(
        peer, CoinState(Coin(bytes32(b"1" * 32), bytes32(b"1" * 32), uint64(0)), uint32(0), uint32(0)), None
    )


@pytest.mark.parametrize(
    "wallet_environments",
    [{"num_environments": 1, "blocks_needed": [1], "trusted": True, "reuse_puzhash": True}],
    indirect=True,
)
@pytest.mark.limit_consensus_modes(reason="irrelevant")
@pytest.mark.anyio
async def test_commit_transactions_to_db(wallet_environments: WalletTestFramework) -> None:
    env = wallet_environments.environments[0]
    wsm = env.wallet_state_manager

    async with wsm.new_action_scope(
        wallet_environments.tx_config,
        push=False,
        merge_spends=False,
        sign=False,
        extra_spends=[],
    ) as action_scope:
        coins = list(await wsm.main_wallet.select_coins(uint64(2_000_000_000_000), action_scope))
        await wsm.main_wallet.generate_signed_transaction(
            [uint64(0)],
            [bytes32.zeros],
            action_scope,
            coins={coins[0]},
        )
        await wsm.main_wallet.generate_signed_transaction(
            [uint64(0)],
            [bytes32.zeros],
            action_scope,
            coins={coins[1]},
        )

    created_txs = action_scope.side_effects.transactions

    def flatten_spend_bundles(txs: list[TransactionRecord]) -> list[WalletSpendBundle]:
        return [tx.spend_bundle for tx in txs if tx.spend_bundle is not None]

    assert (
        len(await wsm.tx_store.get_all_transactions_for_wallet(wsm.main_wallet.id(), type=TransactionType.OUTGOING_TX))
        == 0
    )

    bundles = flatten_spend_bundles(created_txs)
    assert len(bundles) == 2
    for bundle in bundles:
        assert bundle.aggregated_signature == G2Element()
    assert (
        len(await wsm.tx_store.get_all_transactions_for_wallet(wsm.main_wallet.id(), type=TransactionType.OUTGOING_TX))
        == 0
    )

    extra_coin_spend = make_spend(
        Coin(bytes32(b"1" * 32), bytes32(b"1" * 32), uint64(0)), Program.to(1), Program.to([])
    )
    extra_spend = WalletSpendBundle([extra_coin_spend], G2Element())

    new_txs = await wsm.add_pending_transactions(
        created_txs,
        push=False,
        merge_spends=False,
        sign=False,
        extra_spends=[extra_spend],
    )
    bundles = flatten_spend_bundles(new_txs)
    assert len(bundles) == 2
    for bundle in bundles:
        assert bundle.aggregated_signature == G2Element()
    assert (
        len(await wsm.tx_store.get_all_transactions_for_wallet(wsm.main_wallet.id(), type=TransactionType.OUTGOING_TX))
        == 0
    )
    assert extra_coin_spend in [spend for bundle in bundles for spend in bundle.coin_spends]

    new_txs = await wsm.add_pending_transactions(
        created_txs,
        push=False,
        merge_spends=True,
        sign=False,
        extra_spends=[extra_spend],
    )
    bundles = flatten_spend_bundles(new_txs)
    assert len(bundles) == 1
    for bundle in bundles:
        assert bundle.aggregated_signature == G2Element()
    assert (
        len(await wsm.tx_store.get_all_transactions_for_wallet(wsm.main_wallet.id(), type=TransactionType.OUTGOING_TX))
        == 0
    )
    assert extra_coin_spend in [spend for bundle in bundles for spend in bundle.coin_spends]

    new_txs = await wsm.add_pending_transactions(created_txs, push=True, merge_spends=True, sign=True)
    bundles = flatten_spend_bundles(new_txs)
    assert len(bundles) == 1
    assert (
        len(await wsm.tx_store.get_all_transactions_for_wallet(wsm.main_wallet.id(), type=TransactionType.OUTGOING_TX))
        == 2
    )

    await wallet_environments.full_node.wait_transaction_records_entered_mempool(new_txs)
    await wallet_environments.full_node.farm_blocks_to_puzzlehash(count=1, guarantee_transaction_blocks=True)
    await wallet_environments.full_node.wait_for_wallet_synced(wallet_node=env.node, timeout=20)

    rpc_client = env.rpc_client

    spendable_response = await rpc_client.get_spendable_coins(GetSpendableCoins(wallet_id=uint32(1)))
    assert len(spendable_response.confirmed_records) > 0

    select_response = await rpc_client.select_coins(SelectCoins(wallet_id=uint32(1), amount=uint64(1)))
    assert len(select_response.coins) > 0

    spendable = list(await wsm.get_spendable_coins_for_wallet(uint32(1)))
    assert len(spendable) > 0
    coin_name = spendable[0].coin.name()
    records_response = await rpc_client.get_coin_records_by_names(GetCoinRecordsByNames(names=[coin_name]))
    assert len(records_response.coin_records) == 1
    assert records_response.coin_records[0].name == coin_name

    height_response = GetHeightInfoResponse.from_json_dict(
        await rpc_client.fetch("get_height_info", GetHeightInfo().to_json_dict())
    )
    assert height_response.height > 0
    height_peak_response = GetHeightInfoResponse.from_json_dict(
        await rpc_client.fetch("get_height_info", GetHeightInfo(use_peak_height=True).to_json_dict())
    )
    assert height_peak_response.height >= height_response.height


@pytest.mark.parametrize(
    "wallet_environments",
    [{"num_environments": 2, "blocks_needed": [1, 1], "trusted": True, "reuse_puzhash": True}],
    indirect=True,
)
@pytest.mark.limit_consensus_modes(reason="irrelevant")
@pytest.mark.anyio
async def test_confirming_txs_not_ours(wallet_environments: WalletTestFramework) -> None:
    env_1 = wallet_environments.environments[0]
    env_2 = wallet_environments.environments[1]

    # Some transaction, doesn't matter what
    async with env_1.wallet_state_manager.new_action_scope(wallet_environments.tx_config, push=False) as action_scope:
        await env_1.xch_wallet.generate_signed_transaction(
            [uint64(1)],
            [await action_scope.get_puzzle_hash(env_1.wallet_state_manager)],
            action_scope,
        )

    await env_2.rpc_client.push_transactions(
        PushTransactions(
            transactions=action_scope.side_effects.transactions,
            sign=False,
        ),
        wallet_environments.tx_config,
    )

    await wallet_environments.process_pending_states(
        [
            WalletStateTransition(
                pre_block_balance_updates={},
                post_block_balance_updates={
                    1: {
                        "unspent_coin_count": 1,  # We just split a coin so no other balance changes
                    }
                },
            ),
            WalletStateTransition(
                pre_block_balance_updates={
                    1: {
                        "pending_coin_removal_count": 1,  # not sure if this is desirable
                    }
                },
                post_block_balance_updates={
                    1: {
                        "pending_coin_removal_count": -1,
                    }
                },
            ),
        ]
    )


@pytest.mark.parametrize(
    "wallet_environments",
    [{"num_environments": 2, "blocks_needed": [1, 1], "trusted": True, "reuse_puzhash": True}],
    indirect=True,
)
@pytest.mark.limit_consensus_modes(reason="irrelevant")
@pytest.mark.anyio
async def test_confirming_txs_not_ours_with_remote_interest_coin(wallet_environments: WalletTestFramework) -> None:
    """When a REMOTE-only coin record receives a spent update,
    the spent status must be persisted without relying on network calls that
    could fail and roll back the DB transaction.

    Two environments are required because the REMOTE interest-only code path
    in ``_add_coin_states`` is only reachable when ``get_wallet_identifier_for_puzzle_hash``
    returns ``None`` — i.e. the coin's puzzle hash doesn't belong to any wallet
    in this WSM's puzzle store.  If we used a single environment, the standard
    wallet would already claim the puzzle hash, ``wallet_identifier`` would be
    non-None, and the local-record REMOTE fallback (+ subsequent spent short-circuit)
    would never execute.

    env_1 owns the coins and builds the transaction; env_2 only has REMOTE
    interest in the removal coin, guaranteeing the REMOTE path is taken when the
    spent update arrives.
    """
    env_1 = wallet_environments.environments[0]
    env_2 = wallet_environments.environments[1]

    # env_1 builds but does NOT push the tx; env_2 will push it later.
    async with env_1.wallet_state_manager.new_action_scope(wallet_environments.tx_config, push=False) as action_scope:
        await env_1.xch_wallet.generate_signed_transaction(
            [uint64(1)],
            [await action_scope.get_puzzle_hash(env_1.wallet_state_manager)],
            action_scope,
        )

    [tx] = action_scope.side_effects.transactions
    [removed_coin] = tx.removals

    # Register interest in the removal coin to force the REMOTE interest-only branch.
    # Creating via RPC also covers the create_new_wallet endpoint for REMOTE_WALLET.
    response = await env_2.rpc_client.create_new_wallet(
        CreateNewWallet(wallet_type=CreateNewWalletType.REMOTE_WALLET, name="Remote Wallet #1", push=True),
        tx_config=wallet_environments.tx_config,
    )
    assert response.type == WalletType.REMOTE.name
    remote_wallet = env_2.wallet_state_manager.wallets[response.wallet_id]
    assert isinstance(remote_wallet, RemoteWallet)
    removed_coin_id = removed_coin.name()
    await remote_wallet.register_remote_coins([removed_coin_id])

    async def remote_record_spent_flag() -> int:
        record = await env_2.wallet_state_manager.coin_store.get_coin_record(removed_coin_id)
        if record is None:
            return 0
        return int(record.spent)

    # The REMOTE interest record should exist and initially be unspent as coin is not yet spent.
    await time_out_assert(20, remote_record_spent_flag, 0)

    await env_2.rpc_client.push_transactions(
        PushTransactions(
            transactions=action_scope.side_effects.transactions,
            sign=False,
        ),
        wallet_environments.tx_config,
    )

    async def pending_removal_count() -> int:
        balance = (await env_2.rpc_client.get_wallet_balance(GetWalletBalance(wallet_id=uint32(1)))).wallet_balance
        return int(balance.pending_coin_removal_count)

    await time_out_assert(20, pending_removal_count, 1)

    await wallet_environments.full_node.farm_blocks_to_puzzlehash(count=1, guarantee_transaction_blocks=True)
    await wallet_environments.full_node.wait_for_wallet_synced(wallet_node=env_2.node, timeout=20)

    await time_out_assert(20, pending_removal_count, 0)
    # Regression check: spent updates for existing REMOTE records must not be skipped.
    await time_out_assert(20, remote_record_spent_flag, 1)


@dataclass
class PuzzleHashState:
    highest_index: int
    used_up_to_index: int


@pytest.mark.parametrize(
    "wallet_environments",
    [{"num_environments": 1, "blocks_needed": [1], "trusted": True, "reuse_puzhash": True}],
    indirect=True,
)
@pytest.mark.limit_consensus_modes(reason="irrelevant")
@pytest.mark.anyio
async def test_puzzle_hash_requests(wallet_environments: WalletTestFramework) -> None:
    wsm = wallet_environments.environments[0].wallet_state_manager

    async def get_puzzle_hash_state() -> PuzzleHashState:
        last_index = await wsm.puzzle_store.get_last_derivation_path_for_wallet(wsm.main_wallet.id())
        assert last_index is not None
        return PuzzleHashState(
            last_index,
            int((await wsm.puzzle_store.get_used_count(wsm.main_wallet.id())) / 2) - 1,  # hardened + unhardened
        )

    expected_state = await get_puzzle_hash_state()

    # Quick test of this RPC
    assert (
        await wallet_environments.environments[0].rpc_client.get_current_derivation_index()
    ).index == expected_state.highest_index

    # `create_more_puzzle_hashes`
    # No-op
    result = await wsm.create_more_puzzle_hashes()
    await result.commit(wsm)
    assert await get_puzzle_hash_state() == expected_state

    # Ensure the window continues to expand
    await wsm.puzzle_store.set_used_up_to(uint32(expected_state.used_up_to_index + 1))
    result = await wsm.create_more_puzzle_hashes()
    await result.commit(wsm)
    expected_state = PuzzleHashState(expected_state.highest_index + 1, expected_state.used_up_to_index + 1)
    assert await get_puzzle_hash_state() == expected_state

    # Explicitly make 1 extra
    result = await wsm.create_more_puzzle_hashes(num_additional_phs=wsm.initial_num_public_keys + 1)
    await result.commit(wsm)
    expected_state = PuzzleHashState(expected_state.highest_index + 1, expected_state.used_up_to_index)
    assert await get_puzzle_hash_state() == expected_state

    # Make sure window doesn't expand on next use
    await wsm.puzzle_store.set_used_up_to(uint32(expected_state.used_up_to_index + 1))
    result = await wsm.create_more_puzzle_hashes()
    await result.commit(wsm)
    expected_state = PuzzleHashState(expected_state.highest_index, expected_state.used_up_to_index + 1)
    assert await get_puzzle_hash_state() == expected_state

    # Make sure `up_to_index` works
    result = await wsm.create_more_puzzle_hashes(
        up_to_index=uint32(expected_state.highest_index + 100), mark_existing_as_used=False
    )
    await result.commit(wsm)
    expected_state = PuzzleHashState(
        expected_state.highest_index + 100 + wsm.initial_num_public_keys, expected_state.used_up_to_index
    )
    assert await get_puzzle_hash_state() == expected_state

    # Make sure `mark_existing_as_used` works
    result = await wsm.create_more_puzzle_hashes(
        up_to_index=uint32(expected_state.highest_index + 1), mark_existing_as_used=True
    )
    await result.commit(wsm)
    expected_state = PuzzleHashState(
        expected_state.highest_index + 1 + wsm.initial_num_public_keys, expected_state.highest_index
    )
    assert await get_puzzle_hash_state() == expected_state

    # Test basic transactionality
    result = await wsm.create_more_puzzle_hashes(
        up_to_index=uint32(expected_state.highest_index + 1), mark_existing_as_used=False
    )
    result = await wsm.create_more_puzzle_hashes(
        num_additional_phs=(expected_state.highest_index - expected_state.used_up_to_index)
        + wsm.initial_num_public_keys
        + 1,
        mark_existing_as_used=False,
        previous_result=result,
    )
    await result.commit(wsm)
    expected_state = PuzzleHashState(
        expected_state.highest_index + 1 + wsm.initial_num_public_keys + 1, expected_state.used_up_to_index
    )
    assert await get_puzzle_hash_state() == expected_state

    # Test error using two different "config"s
    result = await wsm.create_more_puzzle_hashes(mark_existing_as_used=False)
    with pytest.raises(ValueError, match="different configuration"):
        await wsm.create_more_puzzle_hashes(mark_existing_as_used=True, previous_result=result)

    # Test generation with no local data
    await wsm.puzzle_store.delete_wallet(wsm.main_wallet.id())
    result = await wsm.create_more_puzzle_hashes()
    await result.commit(wsm)
    expected_state = PuzzleHashState(
        wsm.initial_num_public_keys, -1
    )  # -1 being no puzzle hashes used, not even at index 0
    assert await get_puzzle_hash_state() == expected_state

    # Test `from_zero` fills in gaps
    async with wsm.puzzle_store.db_wrapper.writer() as conn:
        await conn.execute(
            "DELETE FROM derivation_paths WHERE derivation_index=?",
            (0,),
        )
    assert await get_puzzle_hash_state() == expected_state
    assert (
        len(list(await wsm.puzzle_store.get_all_puzzle_hashes())) == (expected_state.highest_index) * 2
    )  # 0 inclusive
    result = await wsm.create_more_puzzle_hashes(
        from_zero=True, mark_existing_as_used=False, up_to_index=uint32(expected_state.highest_index)
    )
    await result.commit(wsm)
    expected_state = PuzzleHashState(expected_state.highest_index + wsm.initial_num_public_keys, -1)
    assert await get_puzzle_hash_state() == expected_state
    assert len(list(await wsm.puzzle_store.get_all_puzzle_hashes())) == (expected_state.highest_index + 1) * 2

    # `get_unused_derivation_record`
    # Assert index increases
    assert expected_state.highest_index > expected_state.used_up_to_index
    await wsm.get_unused_derivation_record(wsm.main_wallet.id())
    expected_state = PuzzleHashState(expected_state.highest_index, expected_state.used_up_to_index + 1)
    assert await get_puzzle_hash_state() == expected_state

    # Assert more puzzle hashes get made
    await wsm.puzzle_store.set_used_up_to(uint32(expected_state.highest_index))
    await wsm.get_unused_derivation_record(wsm.main_wallet.id())
    expected_state = PuzzleHashState(
        expected_state.highest_index + wsm.initial_num_public_keys + 1, expected_state.highest_index + 1
    )
    assert await get_puzzle_hash_state() == expected_state

    # Test transactionality
    previous_result = None
    for _ in range(wsm.initial_num_public_keys):  # all currently unused
        previous_result = await wsm._get_unused_derivation_record(wsm.main_wallet.id(), previous_result=previous_result)
    assert previous_result is not None
    await previous_result.commit(wsm)
    expected_state = PuzzleHashState(
        expected_state.highest_index + wsm.initial_num_public_keys, expected_state.highest_index
    )
    assert await get_puzzle_hash_state() == expected_state

    # `extend_derivation_index`
    # Check malformed request
    rpc_client = wallet_environments.environments[0].rpc_client
    with pytest.raises(ValueError):
        await rpc_client.fetch("extend_derivation_index", {})

    # Test no existing derivation paths
    async with wsm.puzzle_store.db_wrapper.writer() as conn:
        await conn.execute(
            "DELETE FROM derivation_paths WHERE derivation_index=?",
            (0,),
        )
    with pytest.raises(ValueError):
        await rpc_client.extend_derivation_index(ExtendDerivationIndex(index=uint32(0)))

    # Reset to a normal state
    await wsm.puzzle_store.delete_wallet(wsm.main_wallet.id())
    result = await wsm.create_more_puzzle_hashes()
    await result.commit(wsm)
    expected_state = PuzzleHashState(wsm.initial_num_public_keys, -1)
    assert await get_puzzle_hash_state() == expected_state

    # Test an index already created
    with pytest.raises(ValueError):
        await rpc_client.extend_derivation_index(ExtendDerivationIndex(index=uint32(0)))

    # Test an index too far in the future
    with pytest.raises(ValueError):
        await rpc_client.extend_derivation_index(
            ExtendDerivationIndex(index=uint32(MAX_DERIVATION_INDEX_DELTA + expected_state.highest_index + 1))
        )

    # Test the actual functionality
    assert (
        await rpc_client.extend_derivation_index(ExtendDerivationIndex(index=uint32(expected_state.highest_index + 5)))
    ).index == expected_state.highest_index + 5
    expected_state = PuzzleHashState(expected_state.highest_index + 5, expected_state.used_up_to_index)
    assert await get_puzzle_hash_state() == expected_state


@pytest.mark.limit_consensus_modes(allowed=[ConsensusMode.HARD_FORK_2_0])
@pytest.mark.anyio
async def test_get_sync_status(simulator_and_wallet: OldSimulatorsAndWallets, self_hostname: str) -> None:
    full_nodes, wallets, _ = simulator_and_wallet
    full_node_api = full_nodes[0]
    full_node_server = full_node_api.full_node.server
    wallet_node, wallet_server = wallets[0]
    await wallet_server.start_client(PeerInfo(self_hostname, full_node_server.get_port()), None)
    wsm: WalletStateManager = wallet_node.wallet_state_manager

    # Farm enough blocks so peak height > 10 (needed for LONG_SYNC test)
    await full_node_api.farm_blocks_to_puzzlehash(count=12, guarantee_transaction_blocks=True)
    await full_node_api.wait_for_wallet_synced(wallet_node=wallet_node, timeout=20)

    original_network = wsm.config["selected_network"]

    # SYNCED via simulator shortcut (line 742)
    wsm.config["selected_network"] = "simulator0"
    try:
        assert await wsm.get_sync_status() == SyncStatus.SYNCED
    finally:
        wsm.config["selected_network"] = original_network

    # SYNCED via normal path (line 754) — wallet is fully synced
    assert await wsm.get_sync_status() == SyncStatus.SYNCED

    # SLIGHTLY_BEHIND (line 752) — height gap small but synced() returns False
    peak = wsm.blockchain._peak
    assert peak is not None
    assert peak.height > 10, f"Need peak > 10 for sync tests, got {peak.height}"
    original_network = wsm.config["selected_network"]
    wsm.config["selected_network"] = "mainnet"
    try:
        await wsm.blockchain.set_finished_sync_up_to(peak.height - 5, in_rollback=True)
        status = await wsm.get_sync_status()
        assert status == SyncStatus.SLIGHTLY_BEHIND, f"Expected SLIGHTLY_BEHIND, got {status}"
    finally:
        wsm.config["selected_network"] = original_network

    # LONG_SYNC (line 749) — set finished_sync_up_to far behind peak
    await wsm.blockchain.set_finished_sync_up_to(0, in_rollback=True)
    assert await wsm.get_sync_status() == SyncStatus.LONG_SYNC

    # DISCONNECTED (line 745) — disconnect from peer
    await wallet_server.close_all_connections()
    await time_out_assert(5, lambda: len(wsm.server.get_connections(NodeType.FULL_NODE)), 0)
    assert await wsm.get_sync_status() == SyncStatus.DISCONNECTED


@pytest.mark.parametrize(
    "wallet_environments",
    [{"num_environments": 1, "blocks_needed": [1], "trusted": True, "reuse_puzhash": True}],
    indirect=True,
)
@pytest.mark.limit_consensus_modes(reason="irrelevant")
@pytest.mark.anyio
async def test_rpc_get_puzzle_and_solution(wallet_environments: WalletTestFramework) -> None:
    env = wallet_environments.environments[0]
    rpc_client = env.rpc_client
    wsm = env.wallet_state_manager
    full_node = wallet_environments.full_node

    async with wsm.new_action_scope(wallet_environments.tx_config, push=True) as action_scope:
        await wsm.main_wallet.generate_signed_transaction(
            [uint64(1)],
            [bytes32.zeros],
            action_scope,
        )

    [tx] = action_scope.side_effects.transactions
    spent_coin = tx.removals[0]

    await full_node.wait_transaction_records_entered_mempool([tx])
    await full_node.farm_blocks_to_puzzlehash(count=1, guarantee_transaction_blocks=True)
    await full_node.wait_for_wallet_synced(wallet_node=env.node, timeout=20)

    request = GetPuzzleAndSolution(coin_name=spent_coin.name()).to_json_dict()
    response = GetPuzzleAndSolutionResponse.from_json_dict(await rpc_client.fetch("get_puzzle_and_solution", request))
    assert response.puzzle_reveal != ""
    assert response.solution != ""

    with pytest.raises(ValueError, match="not found or not spent"):
        await rpc_client.fetch("get_puzzle_and_solution", GetPuzzleAndSolution(coin_name=bytes32.zeros).to_json_dict())


@pytest.mark.parametrize(
    "wallet_environments",
    [{"num_environments": 1, "blocks_needed": [1], "trusted": True, "reuse_puzhash": True}],
    indirect=True,
)
@pytest.mark.limit_consensus_modes(reason="irrelevant")
@pytest.mark.anyio
async def test_rpc_disconnected_errors(wallet_environments: WalletTestFramework) -> None:
    """Verify that RPC endpoints raise when the wallet has no full node peers."""
    env = wallet_environments.environments[0]
    rpc_client = env.rpc_client
    wsm = env.wallet_state_manager

    # Build a dummy PushTransactions request while still connected
    async with wsm.new_action_scope(wallet_environments.tx_config, push=False) as action_scope:
        await wsm.main_wallet.generate_signed_transaction([uint64(1)], [bytes32.zeros], action_scope)
    txs = action_scope.side_effects.transactions

    # Disconnect from all full node peers
    await env.peer_server.close_all_connections()
    await time_out_assert(5, lambda: len(wsm.server.get_connections(NodeType.FULL_NODE)), 0)

    # tx_endpoint decorator DISCONNECTED check (line 341)
    with pytest.raises(ValueError, match="not connected"):
        await rpc_client.push_transactions(
            PushTransactions(transactions=txs, sign=False),
            wallet_environments.tx_config,
        )

    # select_coins DISCONNECTED check
    with pytest.raises(ValueError, match="not connected"):
        await rpc_client.select_coins(SelectCoins(wallet_id=uint32(1), amount=uint64(1)))

    # get_spendable_coins DISCONNECTED check
    with pytest.raises(ValueError, match="not connected"):
        await rpc_client.get_spendable_coins(GetSpendableCoins(wallet_id=uint32(1)))

    # get_coin_records_by_names DISCONNECTED check
    with pytest.raises(ValueError, match="not connected"):
        await rpc_client.get_coin_records_by_names(GetCoinRecordsByNames(names=[bytes32.zeros]))

    # Check the no-peers early exit.
    original_network = wsm.config["selected_network"]
    wsm.config["selected_network"] = "simulator0"
    try:
        with pytest.raises(ValueError, match="No full node peers connected"):
            await rpc_client.get_coin_records_by_names(GetCoinRecordsByNames(names=[bytes32.zeros]))
    finally:
        wsm.config["selected_network"] = original_network


@pytest.mark.parametrize(
    "wallet_environments",
    [{"num_environments": 1, "blocks_needed": [4], "trusted": False, "reuse_puzhash": True}],
    indirect=True,
)
@pytest.mark.limit_consensus_modes(reason="irrelevant")
@pytest.mark.anyio
async def test_get_height_info_with_block_record(wallet_environments: WalletTestFramework) -> None:
    """Ensure the block record branch is hit (untrusted sync populates block records)."""
    env = wallet_environments.environments[0]
    rpc_client = env.rpc_client
    blockchain = env.wallet_state_manager.blockchain

    synced_height = await blockchain.get_finished_sync_up_to()
    assert blockchain.contains_height(synced_height)

    response = GetHeightInfoResponse.from_json_dict(
        await rpc_client.fetch("get_height_info", GetHeightInfo(use_peak_height=False).to_json_dict())
    )
    assert response.height == synced_height
    assert response.is_transaction_block is not None


async def _seed_did_scoped_nft_wallets(wsm: WalletStateManager, did_ids: list[bytes32]) -> list[NFTWallet]:
    """Create one DID-scoped NFT wallet per ``did_ids`` via the real auto-create path."""
    created: list[NFTWallet] = []
    for index, did_id in enumerate(did_ids):
        wallet = await NFTWallet.create_new_nft_wallet(wsm, wsm.main_wallet, did_id=did_id, name=f"NFT {index}")
        created.append(wallet)
    return created


def _build_fake_nft_data(
    *,
    old_p2_puzhash: bytes32,
    singleton_launcher_id: bytes32,
) -> NFTCoinData:
    """Build a duck-typed NFTCoinData for ``handle_nft``.

    ``handle_nft`` only accesses a small subset of fields and the helpers it
    invokes (``get_metadata_and_phs`` and ``get_new_owner_did``) are patched in
    the tests. Constructing real on-chain CoinSpend/UncurriedNFT objects would
    require a full NFT mint, which is orthogonal to the cap behavior under test.
    """
    uncurried_nft = SimpleNamespace(
        supports_did=True,
        owner_did=None,
        p2_puzzle=SimpleNamespace(get_tree_hash=lambda: old_p2_puzhash),
        singleton_launcher_id=singleton_launcher_id,
    )
    parent_coin_spend = SimpleNamespace(
        solution=bytes(Program.to([])),
        coin=SimpleNamespace(),
    )
    parent_coin_state = SimpleNamespace(spent_height=None)
    return cast(
        NFTCoinData,
        SimpleNamespace(
            uncurried_nft=uncurried_nft,
            parent_coin_spend=parent_coin_spend,
            parent_coin_state=parent_coin_state,
        ),
    )


@pytest.mark.limit_consensus_modes(reason="cap logic is consensus-independent")
@pytest.mark.parametrize(
    "case, configured_limit, preexisting_did_ids, seed_matching_wallet, expect_new_wallet, expect_warning",
    [
        pytest.param(
            "at_limit_blocks_creation",
            2,
            [bytes32(b"\x01" * 32), bytes32(b"\x02" * 32)],
            False,
            False,
            True,
            id="at_limit_blocks_creation",
        ),
        pytest.param(
            "below_limit_creates_wallet",
            2,
            [bytes32(b"\x01" * 32)],
            False,
            True,
            False,
            id="below_limit_creates_wallet",
        ),
        pytest.param(
            "matching_wallet_skips_cap",
            1,
            [bytes32(b"\x01" * 32)],
            True,
            False,
            False,
            id="matching_wallet_skips_cap",
        ),
        pytest.param(
            "yaml_default_governs",
            None,
            [],
            False,
            True,
            False,
            id="yaml_default_governs",
        ),
    ],
)
@pytest.mark.anyio
async def test_handle_nft_auto_add_limit(
    simulator_and_wallet: OldSimulatorsAndWallets,
    monkeypatch: pytest.MonkeyPatch,
    caplog: pytest.LogCaptureFixture,
    case: str,
    configured_limit: int | None,
    preexisting_did_ids: list[bytes32],
    seed_matching_wallet: bool,
    expect_new_wallet: bool,
    expect_warning: bool,
) -> None:
    """Regression test for the NFT auto-add cap.

    The inner-puzzle-parsed ``new_did_id`` in NFT transfer data is
    attacker-controllable. Without a cap, ``handle_nft`` would create one
    NFT wallet per unique foreign DID with no upper bound (in contrast to
    the existing ``did_auto_add_limit`` for DID ingestion). This test
    pins down four behaviors:

    - at the configured limit, ``handle_nft`` returns ``None`` and emits a
      warning (no new wallet is created);
    - below the limit, ``handle_nft`` creates a new wallet;
    - when an existing NFT wallet already matches ``new_did_id``, the cap
      is not consulted at all (a regression that hoisted the cap above
      the matching loop would break inbound NFT routing for users at the
      cap);
    - when ``nft_auto_add_limit`` is not present in the config, the
      ``initial-config.yaml`` default of 100 governs — this catches both
      a wrong default literal in the code and any drift between the YAML
      key name and the code's ``config.get`` key.
    """
    _, [(wallet_node, _)], _ = simulator_and_wallet
    wsm = wallet_node.wallet_state_manager

    if configured_limit is not None:
        wsm.config["nft_auto_add_limit"] = configured_limit
    else:
        # The YAML-loaded default must reach the wsm.config dict under the exact
        # key the handle_nft cap reads. Catches both a wrong default literal in
        # initial-config.yaml and any drift between the YAML key and the
        # `config.get("nft_auto_add_limit", ...)` call site in handle_nft.
        assert wsm.config.get("nft_auto_add_limit") == 100

    new_p2_puzhash = bytes32(b"\xaa" * 32)
    old_p2_puzhash = bytes32(b"\xbb" * 32)
    singleton_launcher_id = bytes32(b"\xcc" * 32)
    foreign_did_id = bytes32(b"\xff" * 32)

    sk = master_sk_to_wallet_sk_unhardened(wsm.get_master_private_key(), uint32(99999))
    await wsm.puzzle_store.add_derivation_paths(
        [
            DerivationRecord(
                uint32(99999),
                new_p2_puzhash,
                sk.get_g1(),
                WalletType.STANDARD_WALLET,
                uint32(1),
                False,
            )
        ]
    )

    await _seed_did_scoped_nft_wallets(wsm, preexisting_did_ids)
    if seed_matching_wallet:
        await _seed_did_scoped_nft_wallets(wsm, [foreign_did_id])

    def fake_get_metadata_and_phs(_unft: UncurriedNFT, _solution: bytes) -> tuple[Program, bytes32]:
        return Program.to(0), new_p2_puzhash

    def fake_get_new_owner_did(_unft: UncurriedNFT, _solution: Program) -> bytes32:
        return foreign_did_id

    monkeypatch.setattr(wsm_mod, "get_metadata_and_phs", fake_get_metadata_and_phs)
    monkeypatch.setattr(wsm_mod, "get_new_owner_did", fake_get_new_owner_did)

    nft_data = _build_fake_nft_data(
        old_p2_puzhash=old_p2_puzhash,
        singleton_launcher_id=singleton_launcher_id,
    )

    def nft_wallet_count() -> int:
        return sum(1 for w in wsm.wallets.values() if isinstance(w, NFTWallet))

    before = nft_wallet_count()
    with caplog.at_level(logging.WARNING, logger=wsm.log.name):
        result = await wsm.handle_nft(nft_data)
    after = nft_wallet_count()

    if seed_matching_wallet:
        assert result is not None
        existing_wallet = wsm.wallets[result.id]
        assert isinstance(existing_wallet, NFTWallet)
        assert existing_wallet.nft_wallet_info.did_id == foreign_did_id
        assert after == before
    elif expect_new_wallet:
        assert result is not None
        assert after == before + 1
        new_wallet = wsm.wallets[result.id]
        assert isinstance(new_wallet, NFTWallet)
        assert new_wallet.nft_wallet_info.did_id == foreign_did_id
    else:
        assert result is None
        assert after == before
        assert not any(
            isinstance(w, NFTWallet) and w.nft_wallet_info.did_id == foreign_did_id for w in wsm.wallets.values()
        )

    if expect_warning:
        assert any("nft" in rec.message.lower() and "limit" in rec.message.lower() for rec in caplog.records), (
            f"Expected a cap warning to be emitted; got: {[rec.message for rec in caplog.records]}"
        )
    else:
        assert not any("limit" in rec.message.lower() for rec in caplog.records)