chia-blockchain 2.4.4__py3-none-any.whl

chia/_tests/wallet/test_nft_store.py

@@ -0,0 +1,193 @@
+from __future__ import annotations
+
+import random
+from dataclasses import dataclass, field
+from typing import Dict, List
+
+import pytest
+
+from chia._tests.util.db_connection import DBConnection
+from chia.types.blockchain_format.coin import Coin
+from chia.types.blockchain_format.program import Program
+from chia.types.blockchain_format.sized_bytes import bytes32
+from chia.util.ints import uint32, uint64
+from chia.wallet.lineage_proof import LineageProof
+from chia.wallet.nft_wallet.nft_info import NFTCoinInfo
+from chia.wallet.wallet_nft_store import WalletNftStore
+
+
+def get_dummy_nft(random_: random.Random) -> NFTCoinInfo:
+    return NFTCoinInfo(
+        bytes32.random(random_),
+        Coin(bytes32.random(random_), bytes32.random(random_), uint64(1)),
+        LineageProof(bytes32.random(random_), bytes32.random(random_), uint64(1)),
+        Program.to(["A Test puzzle"]),
+        uint32(1),
+    )
+
+
+@dataclass
+class DummyNFTs:
+    seeded_random: random.Random
+    nfts_per_wallet: Dict[uint32, List[NFTCoinInfo]] = field(default_factory=dict)
+
+    def generate(self, wallet_id: int, count: int) -> None:
+        nfts = self.nfts_per_wallet.setdefault(uint32(wallet_id), [])
+        for _ in range(count):
+            nfts.append(get_dummy_nft(random_=self.seeded_random))
+
+
+class TestNftStore:
+    @pytest.mark.anyio
+    async def test_nft_insert(self) -> None:
+        async with DBConnection(1) as wrapper:
+            db = await WalletNftStore.create(wrapper)
+            a_bytes32 = bytes32.fromhex("09287c75377c63fd6a3a4d6658abed03e9a521e0436b1f83cdf4af99341ce8f1")
+            b_bytes32 = bytes32.fromhex("09287c75377c63fd6a3a4d6658abed03e9a521e0436b1f83cdf4af99341ce8f2")
+            puzzle = Program.to(["A Test puzzle"])
+            nft = NFTCoinInfo(
+                a_bytes32,
+                Coin(a_bytes32, a_bytes32, uint64(1)),
+                LineageProof(a_bytes32, a_bytes32, uint64(1)),
+                puzzle,
+                uint32(1),
+                None,
+                uint32(10),
+            )
+            nft2 = NFTCoinInfo(
+                b_bytes32,
+                Coin(b_bytes32, b_bytes32, uint64(1)),
+                LineageProof(a_bytes32, a_bytes32, uint64(1)),
+                puzzle,
+                uint32(1),
+                None,
+                uint32(10),
+            )
+            assert await db.is_empty()
+            # Test save
+            await db.save_nft(uint32(1), a_bytes32, nft)
+            assert not await db.is_empty()
+            await db.save_nft(uint32(1), a_bytes32, nft)  # test for duplicates
+            await db.save_nft(uint32(1), b_bytes32, nft2)
+            # Test get nft
+            assert await db.count() == 2
+            assert nft == (await db.get_nft_list(wallet_id=uint32(1)))[0]
+            assert nft == (await db.get_nft_list())[0]
+            assert nft == (await db.get_nft_list(did_id=a_bytes32))[0]
+            assert nft == (await db.get_nft_list(wallet_id=uint32(1), did_id=a_bytes32))[0]
+            assert nft == await db.get_nft_by_id(a_bytes32)
+
+            # test get nft list pagination
+            assert not (await db.get_nft_list(wallet_id=uint32(1), start_index=0, count=0))
+            assert nft == (await db.get_nft_list(wallet_id=uint32(1), start_index=0, count=1))[0]
+            assert 1 == len(await db.get_nft_list(wallet_id=uint32(1), start_index=0, count=1))
+            assert 2 == len(await db.get_nft_list(wallet_id=uint32(1), start_index=0, count=5))
+            assert nft2 == (await db.get_nft_list(wallet_id=uint32(1), start_index=0, count=2))[1]
+            assert nft2 == (await db.get_nft_list(wallet_id=uint32(1), start_index=1, count=1))[0]
+            assert 0 == len(await db.get_nft_list(wallet_id=uint32(1), start_index=2, count=1))
+
+            assert nft == (await db.get_nft_by_coin_id(nft.coin.name()))
+            assert await db.exists(nft.coin.name())
+            # negative tests
+            assert (await db.get_nft_by_coin_id(bytes32(b"0" * 32))) is None
+            assert not await db.exists(bytes32(b"0" * 32))
+
+    @pytest.mark.anyio
+    async def test_nft_remove(self) -> None:
+        async with DBConnection(1) as wrapper:
+            db = await WalletNftStore.create(wrapper)
+            a_bytes32 = bytes32.fromhex("09287c75377c63fd6a3a4d6658abed03e9a521e0436b1f83cdf4af99341ce8f1")
+            puzzle = Program.to(["A Test puzzle"])
+            coin = Coin(a_bytes32, a_bytes32, uint64(1))
+            nft = NFTCoinInfo(
+                a_bytes32,
+                coin,
+                LineageProof(a_bytes32, a_bytes32, uint64(1)),
+                puzzle,
+                uint32(1),
+                a_bytes32,
+                uint32(10),
+            )
+            # Test save
+            await db.save_nft(uint32(1), a_bytes32, nft)
+            assert not await db.is_empty()
+            # Test delete by nft id
+            await db.delete_nft_by_nft_id(a_bytes32, uint32(11))
+            assert await db.get_nft_by_id(a_bytes32) is None
+            assert await db.is_empty(wallet_id=uint32(1))
+
+            # Test delete by coin id
+            await db.save_nft(uint32(1), a_bytes32, nft)
+            assert not await db.delete_nft_by_coin_id(a_bytes32, uint32(11))
+            assert await db.delete_nft_by_coin_id(coin.name(), uint32(11))
+            assert not await db.delete_nft_by_coin_id(a_bytes32, uint32(11))
+            assert not await db.exists(a_bytes32)
+
+    @pytest.mark.anyio
+    async def test_nft_reorg(self) -> None:
+        async with DBConnection(1) as wrapper:
+            db = await WalletNftStore.create(wrapper)
+            a_bytes32 = bytes32.fromhex("09287c75377c63fd6a3a4d6658abed03e9a521e0436b1f83cdf4af99341ce8f0")
+            nft_id_1 = bytes32.fromhex("09287c75377c63fd6a3a4d6658abed03e9a521e0436b1f83cdf4af99341ce8f1")
+            coin_id_1 = bytes32.fromhex("09287c75377c63fd6a3a4d6658abed03e9a521e0436b1f83cdf4af99341ce8f2")
+            nft_id_2 = bytes32.fromhex("09287c75377c63fd6a3a4d6658abed03e9a521e0436b1f83cdf4af99341ce8f3")
+            coin_id_2 = bytes32.fromhex("09287c75377c63fd6a3a4d6658abed03e9a521e0436b1f83cdf4af99341ce8f4")
+            puzzle = Program.to(["A Test puzzle"])
+            nft = NFTCoinInfo(
+                nft_id_1,
+                Coin(coin_id_1, coin_id_1, uint64(1)),
+                LineageProof(coin_id_1, coin_id_1, uint64(1)),
+                puzzle,
+                uint32(1),
+                a_bytes32,
+                uint32(10),
+            )
+            # Test save
+            await db.save_nft(uint32(1), nft_id_1, nft)
+            # Test delete
+            await db.delete_nft_by_nft_id(nft_id_1, uint32(11))
+            assert await db.get_nft_by_id(nft_id_1) is None
+            # Test reorg
+            nft1 = NFTCoinInfo(
+                nft_id_2,
+                Coin(coin_id_2, coin_id_2, uint64(1)),
+                LineageProof(coin_id_2, coin_id_2, uint64(1)),
+                puzzle,
+                uint32(1),
+                a_bytes32,
+                uint32(12),
+            )
+            await db.save_nft(uint32(1), nft_id_1, nft1)
+            assert nft1 == (await db.get_nft_list(wallet_id=uint32(1)))[0]
+            assert await db.rollback_to_block(10)
+            assert nft == (await db.get_nft_list(wallet_id=uint32(1)))[0]
+
+            assert not (await db.get_nft_by_coin_id(coin_id_2))
+            assert not (await db.get_nft_by_coin_id(nft_id_1))
+            assert not await db.exists(coin_id_2)
+
+            await db.rollback_to_block(-1)
+            assert await db.count(wallet_id=uint32(1)) == 0
+            assert await db.is_empty(wallet_id=uint32(1))
+
+
+@pytest.mark.anyio
+async def test_delete_wallet(seeded_random: random.Random) -> None:
+    dummy_nfts = DummyNFTs(seeded_random=seeded_random)
+    for i in range(5):
+        dummy_nfts.generate(i, i * 5)
+    async with DBConnection(1) as wrapper:
+        db = await WalletNftStore.create(wrapper)
+        # Add the nfts per wallet and verify them
+        for wallet_id, nfts in dummy_nfts.nfts_per_wallet.items():
+            for nft in nfts:
+                await db.save_nft(wallet_id, None, nft)
+            assert await db.count(wallet_id) == len(nfts)
+        # Remove one wallet after the other and verify before and after each
+        for wallet_id, nfts in dummy_nfts.nfts_per_wallet.items():
+            # Assert the length again here to make sure the previous removals did not affect other wallet_ids
+            assert await db.count(wallet_id) == len(nfts)
+            await db.delete_wallet(wallet_id)
+            assert await db.count(wallet_id) == 0
+
+        assert await db.is_empty()

chia/_tests/wallet/test_notifications.py

@@ -0,0 +1,196 @@
+from __future__ import annotations
+
+import random
+import tempfile
+from pathlib import Path
+from typing import Any, Optional
+
+import pytest
+
+from chia._tests.util.time_out_assert import time_out_assert, time_out_assert_not_none
+from chia.consensus.block_rewards import calculate_base_farmer_reward, calculate_pool_reward
+from chia.simulator.full_node_simulator import FullNodeSimulator
+from chia.simulator.simulator_protocol import FarmNewBlockProtocol
+from chia.types.blockchain_format.sized_bytes import bytes32
+from chia.types.peer_info import PeerInfo
+from chia.util.db_wrapper import DBWrapper2
+from chia.util.ints import uint32, uint64
+from chia.wallet.notification_store import NotificationStore
+from chia.wallet.util.tx_config import DEFAULT_TX_CONFIG
+
+
+# For testing backwards compatibility with a DB change to add height
+@pytest.mark.anyio
+async def test_notification_store_backwards_compat() -> None:
+    # First create the DB the way it would have otheriwse been created
+    with tempfile.TemporaryDirectory() as temporary_directory:
+        db_name = Path(temporary_directory).joinpath("test.sqlite")
+        db_name.parent.mkdir(parents=True, exist_ok=True)
+        async with DBWrapper2.managed(
+            database=db_name,
+        ) as db_wrapper:
+            async with db_wrapper.writer_maybe_transaction() as conn:
+                await conn.execute(
+                    "CREATE TABLE IF NOT EXISTS notifications(coin_id blob PRIMARY KEY,msg blob,amount blob)"
+                )
+                cursor = await conn.execute(
+                    "INSERT OR REPLACE INTO notifications (coin_id, msg, amount) VALUES(?, ?, ?)",
+                    (
+                        bytes32([0] * 32),
+                        bytes([0] * 10),
+                        bytes([0]),
+                    ),
+                )
+                await cursor.close()
+
+            await NotificationStore.create(db_wrapper)
+            await NotificationStore.create(db_wrapper)
+
+
+@pytest.mark.parametrize(
+    "trusted",
+    [True, False],
+)
+@pytest.mark.anyio
+async def test_notifications(
+    self_hostname: str, two_wallet_nodes: Any, trusted: Any, seeded_random: random.Random
+) -> None:
+    full_nodes, wallets, _ = two_wallet_nodes
+    full_node_api: FullNodeSimulator = full_nodes[0]
+    full_node_server = full_node_api.server
+    wallet_node_1, server_0 = wallets[0]
+    wallet_node_2, server_1 = wallets[1]
+    wsm_1 = wallet_node_1.wallet_state_manager
+    wsm_2 = wallet_node_2.wallet_state_manager
+    wallet_1 = wsm_1.main_wallet
+    wallet_2 = wsm_2.main_wallet
+
+    ph_1 = await wallet_1.get_new_puzzlehash()
+    ph_2 = await wallet_2.get_new_puzzlehash()
+    ph_token = bytes32.random(seeded_random)
+
+    if trusted:
+        wallet_node_1.config["trusted_peers"] = {
+            full_node_api.full_node.server.node_id.hex(): full_node_api.full_node.server.node_id.hex()
+        }
+        wallet_node_2.config["trusted_peers"] = {
+            full_node_api.full_node.server.node_id.hex(): full_node_api.full_node.server.node_id.hex()
+        }
+    else:
+        wallet_node_1.config["trusted_peers"] = {}
+        wallet_node_2.config["trusted_peers"] = {}
+
+    await server_0.start_client(PeerInfo(self_hostname, full_node_server.get_port()), None)
+    await server_1.start_client(PeerInfo(self_hostname, full_node_server.get_port()), None)
+
+    for i in range(0, 2):
+        await full_node_api.farm_new_transaction_block(FarmNewBlockProtocol(ph_1))
+    await full_node_api.farm_new_transaction_block(FarmNewBlockProtocol(ph_token))
+    await full_node_api.wait_for_wallets_synced(wallet_nodes=[wallet_node_1, wallet_node_2], timeout=30)
+
+    funds_1 = sum(calculate_pool_reward(uint32(i)) + calculate_base_farmer_reward(uint32(i)) for i in range(1, 3))
+    funds_2 = 0
+
+    await time_out_assert(30, wallet_1.get_unconfirmed_balance, funds_1)
+    await time_out_assert(30, wallet_1.get_confirmed_balance, funds_1)
+
+    notification_manager_1 = wsm_1.notification_manager
+    notification_manager_2 = wsm_2.notification_manager
+
+    func = notification_manager_2.potentially_add_new_notification
+    notification_manager_2.most_recent_args = tuple()
+
+    async def track_coin_state(*args: Any) -> bool:
+        notification_manager_2.most_recent_args = args
+        result: bool = await func(*args)
+        return result
+
+    notification_manager_2.potentially_add_new_notification = track_coin_state
+
+    allow_larger_height: Optional[int] = None
+    allow_height: Optional[int] = None
+
+    for case in ("block all", "block too low", "allow", "allow_larger", "block_too_large"):
+        msg: bytes = bytes(case, "utf8")
+        if case == "block all":
+            wallet_node_2.config["enable_notifications"] = False
+            wallet_node_2.config["required_notification_amount"] = 100
+            AMOUNT = uint64(100)
+            FEE = uint64(0)
+        elif case == "block too low":
+            wallet_node_2.config["enable_notifications"] = True
+            AMOUNT = uint64(1)
+            FEE = uint64(0)
+        elif case in ("allow", "allow_larger"):
+            wallet_node_2.config["required_notification_amount"] = 750000000000
+            if case == "allow_larger":
+                AMOUNT = uint64(1000000000000)
+            else:
+                AMOUNT = uint64(750000000000)
+            FEE = uint64(1)
+        elif case == "block_too_large":
+            msg = bytes([0] * 10001)
+            AMOUNT = uint64(750000000000)
+            FEE = uint64(0)
+        else:
+            raise Exception(f"Unhandled case: {case!r}")
+        peak = full_node_api.full_node.blockchain.get_peak()
+        assert peak is not None
+        if case == "allow":
+            allow_height = peak.height + 1
+        if case == "allow_larger":
+            allow_larger_height = peak.height + 1
+        async with notification_manager_1.wallet_state_manager.new_action_scope(
+            DEFAULT_TX_CONFIG, push=True
+        ) as action_scope:
+            await notification_manager_1.send_new_notification(ph_2, msg, AMOUNT, action_scope, fee=FEE)
+        [tx] = action_scope.side_effects.transactions
+        await time_out_assert_not_none(
+            5,
+            full_node_api.full_node.mempool_manager.get_spendbundle,
+            tx.spend_bundle.name(),
+        )
+        await full_node_api.farm_new_transaction_block(FarmNewBlockProtocol(ph_token))
+
+        funds_1 = funds_1 - AMOUNT - FEE
+        funds_2 += AMOUNT
+
+        await time_out_assert(30, wallet_1.get_unconfirmed_balance, funds_1)
+        await time_out_assert(30, wallet_1.get_confirmed_balance, funds_1)
+        await time_out_assert(30, wallet_2.get_unconfirmed_balance, funds_2)
+        await time_out_assert(30, wallet_2.get_confirmed_balance, funds_2)
+
+    assert allow_larger_height is not None
+    assert allow_height is not None
+
+    notifications = await notification_manager_2.notification_store.get_all_notifications(pagination=(0, 2))
+    assert len(notifications) == 2
+    assert notifications[0].message == b"allow_larger"
+    assert notifications[0].height == allow_larger_height
+    notifications = await notification_manager_2.notification_store.get_all_notifications(pagination=(1, None))
+    assert len(notifications) == 1
+    assert notifications[0].message == b"allow"
+    assert notifications[0].height == allow_height
+    notifications = await notification_manager_2.notification_store.get_all_notifications(pagination=(0, 1))
+    assert len(notifications) == 1
+    assert notifications[0].message == b"allow_larger"
+    notifications = await notification_manager_2.notification_store.get_all_notifications(pagination=(None, 1))
+    assert len(notifications) == 1
+    assert notifications[0].message == b"allow_larger"
+    assert (
+        await notification_manager_2.notification_store.get_notifications([n.id for n in notifications])
+        == notifications
+    )
+
+    sent_notifications = await notification_manager_1.notification_store.get_all_notifications()
+    assert len(sent_notifications) == 0
+
+    await notification_manager_2.notification_store.delete_all_notifications()
+    assert len(await notification_manager_2.notification_store.get_all_notifications()) == 0
+    await notification_manager_2.notification_store.add_notification(notifications[0])
+    await notification_manager_2.notification_store.delete_notifications([n.id for n in notifications])
+    assert len(await notification_manager_2.notification_store.get_all_notifications()) == 0
+
+    assert not await func(*notification_manager_2.most_recent_args)
+    await notification_manager_2.notification_store.delete_all_notifications()
+    assert not await func(*notification_manager_2.most_recent_args)

chia/_tests/wallet/test_offer_parsing_performance.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+import cProfile
+from contextlib import contextmanager
+from typing import Iterator
+
+from chia._tests.util.misc import BenchmarkRunner
+from chia.wallet.trading.offer import Offer
+
+with_profile = False
+
+# gprof2dot -f pstats offer-parsing.profile >p.dot && dot -Tpng p.dot >offer-parsing.png
+# gprof2dot -f pstats offered-coins.profile >c.dot && dot -Tpng c.dot >offered-coins.png
+
+
+@contextmanager
+def enable_profiler(name: str) -> Iterator[None]:
+    if not with_profile:
+        yield
+        return
+
+    with cProfile.Profile() as pr:
+        yield
+
+    pr.create_stats()
+    pr.dump_stats(f"{name}.profile")
+
+
+def test_offer_parsing_performance(benchmark_runner: BenchmarkRunner) -> None:
+    offer_bytes = bytes.fromhex(test_offer)
+    with benchmark_runner.assert_runtime(seconds=2.15):
+        with enable_profiler("offer-parsing"):
+            for _ in range(100):
+                o = Offer.from_bytes(offer_bytes)
+                assert o is not None
+
+
+def test_offered_coins_performance(benchmark_runner: BenchmarkRunner) -> None:
+    offer_bytes = bytes.fromhex(test_offer)
+    o = Offer.from_bytes(offer_bytes)
+    with benchmark_runner.assert_runtime(seconds=2.5):
+        with enable_profiler("offered-coins"):
+            for _ in range(100):
+                c = o.get_offered_coins()
+                assert len(c.items()) > 0
+
+
+test_offer = "000000050000000000000000000000000000000000000000000000000000000000000000131816d2ed127bfc4fcde644a6f46e764b9123a9d041b22726d35ab5ff13a19c0000000000000000ff02ffff01ff02ffff01ff02ff5effff04ff02ffff04ffff04ff05ffff04ffff0bff34ff0580ffff04ff0bff80808080ffff04ffff02ff17ff2f80ffff04ff5fffff04ffff02ff2effff04ff02ffff04ff17ff80808080ffff04ffff02ff2affff04ff02ffff04ff82027fffff04ff82057fffff04ff820b7fff808080808080ffff04ff81bfffff04ff82017fffff04ff8202ffffff04ff8205ffffff04ff820bffff80808080808080808080808080ffff04ffff01ffffffff3d46ff02ff333cffff0401ff01ff81cb02ffffff20ff02ffff03ff05ffff01ff02ff32ffff04ff02ffff04ff0dffff04ffff0bff7cffff0bff34ff2480ffff0bff7cffff0bff7cffff0bff34ff2c80ff0980ffff0bff7cff0bffff0bff34ff8080808080ff8080808080ffff010b80ff0180ffff02ffff03ffff22ffff09ffff0dff0580ff2280ffff09ffff0dff0b80ff2280ffff15ff17ffff0181ff8080ffff01ff0bff05ff0bff1780ffff01ff088080ff0180ffff02ffff03ff0bffff01ff02ffff03ffff09ffff02ff2effff04ff02ffff04ff13ff80808080ff820b9f80ffff01ff02ff56ffff04ff02ffff04ffff02ff13ffff04ff5fffff04ff17ffff04ff2fffff04ff81bfffff04ff82017fffff04ff1bff8080808080808080ffff04ff82017fff8080808080ffff01ff088080ff0180ffff01ff02ffff03ff17ffff01ff02ffff03ffff20ff81bf80ffff0182017fffff01ff088080ff0180ffff01ff088080ff018080ff0180ff04ffff04ff05ff2780ffff04ffff10ff0bff5780ff778080ffffff02ffff03ff05ffff01ff02ffff03ffff09ffff02ffff03ffff09ff11ff5880ffff0159ff8080ff0180ffff01818f80ffff01ff02ff26ffff04ff02ffff04ff0dffff04ff0bffff04ffff04ff81b9ff82017980ff808080808080ffff01ff02ff7affff04ff02ffff04ffff02ffff03ffff09ff11ff5880ffff01ff04ff58ffff04ffff02ff76ffff04ff02ffff04ff13ffff04ff29ffff04ffff0bff34ff5b80ffff04ff2bff80808080808080ff398080ffff01ff02ffff03ffff09ff11ff7880ffff01ff02ffff03ffff20ffff02ffff03ffff09ffff0121ffff0dff298080ffff01ff02ffff03ffff09ffff0cff29ff80ff3480ff5c80ffff01ff0101ff8080ff0180ff8080ff018080ffff0109ffff01ff088080ff0180ffff010980ff018080ff0180ffff04ffff02ffff03ffff09ff11ff5880ffff0159ff8080ff0180ffff04ffff02ff26ffff04ff02ffff04ff0dffff04ff0bffff04ff17ff808080808080ff80808080808080ff0180ffff01ff04ff80ffff04ff80ff17808080ff0180ffff02ffff03ff05ffff01ff04ff09ffff02ff56ffff04ff02ffff04ff0dffff04ff0bff808080808080ffff010b80ff0180ff0bff7cffff0bff34ff2880ffff0bff7cffff0bff7cffff0bff34ff2c80ff0580ffff0bff7cffff02ff32ffff04ff02ffff04ff07ffff04ffff0bff34ff3480ff8080808080ffff0bff34ff8080808080ffff02ffff03ffff07ff0580ffff01ff0bffff0102ffff02ff2effff04ff02ffff04ff09ff80808080ffff02ff2effff04ff02ffff04ff0dff8080808080ffff01ff0bffff0101ff058080ff0180ffff04ffff04ff30ffff04ff5fff808080ffff02ff7effff04ff02ffff04ffff04ffff04ff2fff0580ffff04ff5fff82017f8080ffff04ffff02ff26ffff04ff02ffff04ff0bffff04ff05ffff01ff808080808080ffff04ff17ffff04ff81bfffff04ff82017fffff04ffff02ff2affff04ff02ffff04ff8204ffffff04ffff02ff76ffff04ff02ffff04ff09ffff04ff820affffff04ffff0bff34ff2d80ffff04ff15ff80808080808080ffff04ff8216ffff808080808080ffff04ff8205ffffff04ff820bffff808080808080808080808080ff02ff5affff04ff02ffff04ff5fffff04ff3bffff04ffff02ffff03ff17ffff01ff09ff2dffff02ff2affff04ff02ffff04ff27ffff04ffff02ff76ffff04ff02ffff04ff29ffff04ff57ffff04ffff0bff34ff81b980ffff04ff59ff80808080808080ffff04ff81b7ff80808080808080ff8080ff0180ffff04ff17ffff04ff05ffff04ff8202ffffff04ffff04ffff04ff78ffff04ffff0eff5cffff02ff2effff04ff02ffff04ffff04ff2fffff04ff82017fff808080ff8080808080ff808080ffff04ffff04ff20ffff04ffff0bff81bfff5cffff02ff2effff04ff02ffff04ffff04ff15ffff04ffff10ff82017fffff11ff8202dfff2b80ff8202ff80ff808080ff8080808080ff808080ff138080ff80808080808080808080ff018080ffff04ffff01a037bef360ee858133b69d595a906dc45d01af50379dad515eb9518abb7c1d2a7affff04ffff01a0b2847c2a1428be2af434b78ab85b08c46f13cb4dfd686ac21e7ee484504bebacffff04ffff01ff02ffff01ff02ff0affff04ff02ffff04ff03ff80808080ffff04ffff01ffff333effff02ffff03ff05ffff01ff04ffff04ff0cffff04ffff02ff1effff04ff02ffff04ff09ff80808080ff808080ffff02ff16ffff04ff02ffff04ff19ffff04ffff02ff0affff04ff02ffff04ff0dff80808080ff808080808080ff8080ff0180ffff02ffff03ff05ffff01ff02ffff03ffff15ff29ff8080ffff01ff04ffff04ff08ff0980ffff02ff16ffff04ff02ffff04ff0dffff04ff0bff808080808080ffff01ff088080ff0180ffff010b80ff0180ff02ffff03ffff07ff0580ffff01ff0bffff0102ffff02ff1effff04ff02ffff04ff09ff80808080ffff02ff1effff04ff02ffff04ff0dff8080808080ffff01ff0bffff0101ff058080ff0180ff018080ff0180808080ffffa05dd3364e78b4362d07eb167bd9777a27198ea20f2d365a8b99d1f7ee9480678cffffa0e4fb5940aef16e2c6cfed28ea3934dada96f4ce2b629b17cd482eb31f6a6c559ff02ffffa0e4fb5940aef16e2c6cfed28ea3934dada96f4ce2b629b17cd482eb31f6a6c559808080800000000000000000000000000000000000000000000000000000000000000000cfbfdeed5c4ca2de3d0bf520b9cb4bb7743a359bd2e6a188d19ce7dffc21d3e70000000000000000ff02ffff01ff02ff0affff04ff02ffff04ff03ff80808080ffff04ffff01ffff333effff02ffff03ff05ffff01ff04ffff04ff0cffff04ffff02ff1effff04ff02ffff04ff09ff80808080ff808080ffff02ff16ffff04ff02ffff04ff19ffff04ffff02ff0affff04ff02ffff04ff0dff80808080ff808080808080ff8080ff0180ffff02ffff03ff05ffff01ff02ffff03ffff15ff29ff8080ffff01ff04ffff04ff08ff0980ffff02ff16ffff04ff02ffff04ff0dffff04ff0bff808080808080ffff01ff088080ff0180ffff010b80ff0180ff02ffff03ffff07ff0580ffff01ff0bffff0102ffff02ff1effff04ff02ffff04ff09ff80808080ffff02ff1effff04ff02ffff04ff0dff8080808080ffff01ff0bffff0101ff058080ff0180ff018080ffffa0029457911f592c7b41905e7ff40c6db95307b0da9ea68c44c1714d2c5b7ff857ffffa0ba934771c8d2d08140a6bbe9a7508f192b64564bf30a054e3cf215224845dd5aff01ff808080809563629e653a9fc3c65f55947883a47e062e6b67394091228ec01352ff78f333e4fb5940aef16e2c6cfed28ea3934dada96f4ce2b629b17cd482eb31f6a6c5590000003a3529439cff02ffff01ff02ffff01ff02ffff03ff0bffff01ff02ffff03ffff09ff05ffff1dff0bffff1effff0bff0bffff02ff06ffff04ff02ffff04ff17ff8080808080808080ffff01ff02ff17ff2f80ffff01ff088080ff0180ffff01ff04ffff04ff04ffff04ff05ffff04ffff02ff06ffff04ff02ffff04ff17ff80808080ff80808080ffff02ff17ff2f808080ff0180ffff04ffff01ff32ff02ffff03ffff07ff0580ffff01ff0bffff0102ffff02ff06ffff04ff02ffff04ff09ff80808080ffff02ff06ffff04ff02ffff04ff0dff8080808080ffff01ff0bffff0101ff058080ff0180ff018080ffff04ffff01b080399dec21ec7ab66ed0b43811652c2c27c16e2f304b8e2900d7f510448675ef9d8be70ba19e6b6675c4d9bb5f75ced4ff018080ff80ffff01ffff33ffa0cfbfdeed5c4ca2de3d0bf520b9cb4bb7743a359bd2e6a188d19ce7dffc21d3e7ff0180ffff33ffa0e4fb5940aef16e2c6cfed28ea3934dada96f4ce2b629b17cd482eb31f6a6c559ff853a3529439a80ffff34ff0180ffff3cffa0c894876024c3dfa16dbe7cd8ae7c3f51315c6b094d67e56fd043e6dfe35deb8180ffff3fffa08f91df7c4d24e94cfbf9de74dda2dc8795ed52c66e22b4ba9bb7be6a8241432b8080ff80804b3dbe3376b29f3d14062ae7450c27448902dbd7d01eee959d8e1a3725ac70e4240131d2d424812c8a4a9756209100da8b9f80081e098053b4e5348d13fbfaf40000000000000064ff02ffff01ff02ffff01ff02ff5effff04ff02ffff04ffff04ff05ffff04ffff0bff34ff0580ffff04ff0bff80808080ffff04ffff02ff17ff2f80ffff04ff5fffff04ffff02ff2effff04ff02ffff04ff17ff80808080ffff04ffff02ff2affff04ff02ffff04ff82027fffff04ff82057fffff04ff820b7fff808080808080ffff04ff81bfffff04ff82017fffff04ff8202ffffff04ff8205ffffff04ff820bffff80808080808080808080808080ffff04ffff01ffffffff3d46ff02ff333cffff0401ff01ff81cb02ffffff20ff02ffff03ff05ffff01ff02ff32ffff04ff02ffff04ff0dffff04ffff0bff7cffff0bff34ff2480ffff0bff7cffff0bff7cffff0bff34ff2c80ff0980ffff0bff7cff0bffff0bff34ff8080808080ff8080808080ffff010b80ff0180ffff02ffff03ffff22ffff09ffff0dff0580ff2280ffff09ffff0dff0b80ff2280ffff15ff17ffff0181ff8080ffff01ff0bff05ff0bff1780ffff01ff088080ff0180ffff02ffff03ff0bffff01ff02ffff03ffff09ffff02ff2effff04ff02ffff04ff13ff80808080ff820b9f80ffff01ff02ff56ffff04ff02ffff04ffff02ff13ffff04ff5fffff04ff17ffff04ff2fffff04ff81bfffff04ff82017fffff04ff1bff8080808080808080ffff04ff82017fff8080808080ffff01ff088080ff0180ffff01ff02ffff03ff17ffff01ff02ffff03ffff20ff81bf80ffff0182017fffff01ff088080ff0180ffff01ff088080ff018080ff0180ff04ffff04ff05ff2780ffff04ffff10ff0bff5780ff778080ffffff02ffff03ff05ffff01ff02ffff03ffff09ffff02ffff03ffff09ff11ff5880ffff0159ff8080ff0180ffff01818f80ffff01ff02ff26ffff04ff02ffff04ff0dffff04ff0bffff04ffff04ff81b9ff82017980ff808080808080ffff01ff02ff7affff04ff02ffff04ffff02ffff03ffff09ff11ff5880ffff01ff04ff58ffff04ffff02ff76ffff04ff02ffff04ff13ffff04ff29ffff04ffff0bff34ff5b80ffff04ff2bff80808080808080ff398080ffff01ff02ffff03ffff09ff11ff7880ffff01ff02ffff03ffff20ffff02ffff03ffff09ffff0121ffff0dff298080ffff01ff02ffff03ffff09ffff0cff29ff80ff3480ff5c80ffff01ff0101ff8080ff0180ff8080ff018080ffff0109ffff01ff088080ff0180ffff010980ff018080ff0180ffff04ffff02ffff03ffff09ff11ff5880ffff0159ff8080ff0180ffff04ffff02ff26ffff04ff02ffff04ff0dffff04ff0bffff04ff17ff808080808080ff80808080808080ff0180ffff01ff04ff80ffff04ff80ff17808080ff0180ffff02ffff03ff05ffff01ff04ff09ffff02ff56ffff04ff02ffff04ff0dffff04ff0bff808080808080ffff010b80ff0180ff0bff7cffff0bff34ff2880ffff0bff7cffff0bff7cffff0bff34ff2c80ff0580ffff0bff7cffff02ff32ffff04ff02ffff04ff07ffff04ffff0bff34ff3480ff8080808080ffff0bff34ff8080808080ffff02ffff03ffff07ff0580ffff01ff0bffff0102ffff02ff2effff04ff02ffff04ff09ff80808080ffff02ff2effff04ff02ffff04ff0dff8080808080ffff01ff0bffff0101ff058080ff0180ffff04ffff04ff30ffff04ff5fff808080ffff02ff7effff04ff02ffff04ffff04ffff04ff2fff0580ffff04ff5fff82017f8080ffff04ffff02ff26ffff04ff02ffff04ff0bffff04ff05ffff01ff808080808080ffff04ff17ffff04ff81bfffff04ff82017fffff04ffff02ff2affff04ff02ffff04ff8204ffffff04ffff02ff76ffff04ff02ffff04ff09ffff04ff820affffff04ffff0bff34ff2d80ffff04ff15ff80808080808080ffff04ff8216ffff808080808080ffff04ff8205ffffff04ff820bffff808080808080808080808080ff02ff5affff04ff02ffff04ff5fffff04ff3bffff04ffff02ffff03ff17ffff01ff09ff2dffff02ff2affff04ff02ffff04ff27ffff04ffff02ff76ffff04ff02ffff04ff29ffff04ff57ffff04ffff0bff34ff81b980ffff04ff59ff80808080808080ffff04ff81b7ff80808080808080ff8080ff0180ffff04ff17ffff04ff05ffff04ff8202ffffff04ffff04ffff04ff78ffff04ffff0eff5cffff02ff2effff04ff02ffff04ffff04ff2fffff04ff82017fff808080ff8080808080ff808080ffff04ffff04ff20ffff04ffff0bff81bfff5cffff02ff2effff04ff02ffff04ffff04ff15ffff04ffff10ff82017fffff11ff8202dfff2b80ff8202ff80ff808080ff8080808080ff808080ff138080ff80808080808080808080ff018080ffff04ffff01a037bef360ee858133b69d595a906dc45d01af50379dad515eb9518abb7c1d2a7affff04ffff01a0b2847c2a1428be2af434b78ab85b08c46f13cb4dfd686ac21e7ee484504bebacffff04ffff01ff02ffff01ff02ffff01ff02ffff03ff0bffff01ff02ffff03ffff09ff05ffff1dff0bffff1effff0bff0bffff02ff06ffff04ff02ffff04ff17ff8080808080808080ffff01ff02ff17ff2f80ffff01ff088080ff0180ffff01ff04ffff04ff04ffff04ff05ffff04ffff02ff06ffff04ff02ffff04ff17ff80808080ff80808080ffff02ff17ff2f808080ff0180ffff04ffff01ff32ff02ffff03ffff07ff0580ffff01ff0bffff0102ffff02ff06ffff04ff02ffff04ff09ff80808080ffff02ff06ffff04ff02ffff04ff0dff8080808080ffff01ff0bffff0101ff058080ff0180ff018080ffff04ffff01b0a6c878d82fbd95f75e4ddfddb585d356ca350c66abd7e7fecd8de3742d0bef45bae59ac94285eef0bafa91941900a5d2ff018080ff0180808080ffff80ffff01ffff33ffa0cfbfdeed5c4ca2de3d0bf520b9cb4bb7743a359bd2e6a188d19ce7dffc21d3e7ff02ffffa0cfbfdeed5c4ca2de3d0bf520b9cb4bb7743a359bd2e6a188d19ce7dffc21d3e78080ffff33ffa00532eb4ea071f8f965a9269ed6db179daf5f3799a27b0f9d218bfced3d58bffbff62ffffa00532eb4ea071f8f965a9269ed6db179daf5f3799a27b0f9d218bfced3d58bffb8080ffff3cffa0253d152f54d761fc1760b443edb0dcf979d4b8167d2c2e4aada4685657942fa480ffff3fffa015cdba2758ee18412ade437dd0cf668516b3154184ba84121ef3de1dc1f848798080ff8080ffffa0c8b24907a9f259144c45a36fc61177496d272167db2428a1b59e6a248597086cffa00532eb4ea071f8f965a9269ed6db179daf5f3799a27b0f9d218bfced3d58bffbff6480ffa050687dacfe142b76db5af3bc3bd197bd89c0287e7665e5925573fd4e672b23c5ffffa04b3dbe3376b29f3d14062ae7450c27448902dbd7d01eee959d8e1a3725ac70e4ffa0240131d2d424812c8a4a9756209100da8b9f80081e098053b4e5348d13fbfaf4ff6480ffffa04b3dbe3376b29f3d14062ae7450c27448902dbd7d01eee959d8e1a3725ac70e4ffa00532eb4ea071f8f965a9269ed6db179daf5f3799a27b0f9d218bfced3d58bffbff6480ff80ff8080c8b24907a9f259144c45a36fc61177496d272167db2428a1b59e6a248597086cba934771c8d2d08140a6bbe9a7508f192b64564bf30a054e3cf215224845dd5a0000003a3529439cff02ffff01ff02ffff01ff02ffff03ff0bffff01ff02ffff03ffff09ff05ffff1dff0bffff1effff0bff0bffff02ff06ffff04ff02ffff04ff17ff8080808080808080ffff01ff02ff17ff2f80ffff01ff088080ff0180ffff01ff04ffff04ff04ffff04ff05ffff04ffff02ff06ffff04ff02ffff04ff17ff80808080ff80808080ffff02ff17ff2f808080ff0180ffff04ffff01ff32ff02ffff03ffff07ff0580ffff01ff0bffff0102ffff02ff06ffff04ff02ffff04ff09ff80808080ffff02ff06ffff04ff02ffff04ff0dff8080808080ffff01ff0bffff0101ff058080ff0180ff018080ffff04ffff01b08595a83f106ce887c37b956bf828b8aafc677bff1b07e436843eb1201dde9120d0420634cb362ac2bcf8de6454917bf2ff018080ff80ffff01ffff33ffa0ba934771c8d2d08140a6bbe9a7508f192b64564bf30a054e3cf215224845dd5aff853a3529439b80ffff34ff0180ffff3cffa03c0aa0950679baa320f83fabf466a7261a10ff6c173caea22fb4e4afa470ee7e80ffff3dffa03f89b02b428b0199ba475cbf9e9e66eaaa6d751c424d720349e82f8d3cfba7258080ff8080b1df9543cfd23284a1b9967156289c39d6f3dffde041eb185cf37f0573d50f1b0b5b8efd3e7c63a210e8141c5e9ad4cd0124d360a05c44e92b454df6335c404e87edfbf1e85e79cc569d8c44a5fa53eb6f4092ec51ebc99d67a991dd92742311"  # noqa: E501

chia/_tests/wallet/test_puzzle_store.py

@@ -0,0 +1,133 @@
+from __future__ import annotations
+
+import random
+from dataclasses import dataclass, field
+from typing import Dict, List
+
+import pytest
+from chia_rs import AugSchemeMPL
+
+from chia._tests.util.db_connection import DBConnection
+from chia.types.blockchain_format.sized_bytes import bytes32
+from chia.util.ints import uint32
+from chia.wallet.derivation_record import DerivationRecord
+from chia.wallet.util.wallet_types import WalletIdentifier, WalletType
+from chia.wallet.wallet_puzzle_store import WalletPuzzleStore
+
+
+def get_dummy_record(index: int, wallet_id: int, seeded_random: random.Random) -> DerivationRecord:
+    return DerivationRecord(
+        uint32(index),
+        bytes32.random(seeded_random),
+        AugSchemeMPL.key_gen(bytes32.random(seeded_random)).get_g1(),
+        WalletType.STANDARD_WALLET,
+        uint32(wallet_id),
+        False,
+    )
+
+
+@dataclass
+class DummyDerivationRecords:
+    seeded_random: random.Random
+    index_per_wallet: Dict[int, int] = field(default_factory=dict)
+    records_per_wallet: Dict[int, List[DerivationRecord]] = field(default_factory=dict)
+
+    def generate(self, wallet_id: int, count: int) -> None:
+        records = self.records_per_wallet.setdefault(wallet_id, [])
+        self.index_per_wallet.setdefault(wallet_id, 0)
+        for _ in range(count):
+            records.append(
+                get_dummy_record(self.index_per_wallet[wallet_id], wallet_id, seeded_random=self.seeded_random)
+            )
+            self.index_per_wallet[wallet_id] += 1
+
+
+@pytest.mark.anyio
+async def test_puzzle_store(seeded_random: random.Random) -> None:
+    async with DBConnection(1) as wrapper:
+        db = await WalletPuzzleStore.create(wrapper)
+        derivation_recs = []
+        for i in range(1000):
+            derivation_recs.append(
+                DerivationRecord(
+                    uint32(i),
+                    bytes32.random(seeded_random),
+                    AugSchemeMPL.key_gen(bytes32.random(seeded_random)).get_g1(),
+                    WalletType.STANDARD_WALLET,
+                    uint32(1),
+                    False,
+                )
+            )
+            derivation_recs.append(
+                DerivationRecord(
+                    uint32(i),
+                    bytes32.random(seeded_random),
+                    AugSchemeMPL.key_gen(bytes32.random(seeded_random)).get_g1(),
+                    WalletType.CAT,
+                    uint32(2),
+                    False,
+                )
+            )
+        assert await db.puzzle_hash_exists(derivation_recs[0].puzzle_hash) is False
+        assert await db.index_for_pubkey(derivation_recs[0].pubkey) is None
+        assert await db.index_for_puzzle_hash(derivation_recs[2].puzzle_hash) is None
+        assert await db.get_wallet_identifier_for_puzzle_hash(derivation_recs[2].puzzle_hash) is None
+        assert len(await db.get_all_puzzle_hashes()) == 0
+        assert await db.get_last_derivation_path() is None
+        assert await db.get_unused_derivation_path() is None
+        assert await db.get_derivation_record(0, 2, False) is None
+
+        await db.add_derivation_paths(derivation_recs)
+
+        assert await db.puzzle_hash_exists(derivation_recs[0].puzzle_hash) is True
+
+        assert await db.index_for_pubkey(derivation_recs[4].pubkey) == 2
+        assert await db.index_for_puzzle_hash(derivation_recs[2].puzzle_hash) == 1
+        assert await db.get_wallet_identifier_for_puzzle_hash(derivation_recs[2].puzzle_hash) == WalletIdentifier(
+            derivation_recs[2].wallet_id,
+            derivation_recs[2].wallet_type,
+        )
+        assert len(await db.get_all_puzzle_hashes()) == 2000
+        assert await db.get_last_derivation_path() == 999
+        assert await db.get_unused_derivation_path() == 0
+        assert await db.get_derivation_record(0, 2, False) == derivation_recs[1]
+
+        # Indeces up to 250
+        await db.set_used_up_to(249)
+
+        assert await db.get_unused_derivation_path() == 250
+
+
+@pytest.mark.anyio
+async def test_delete_wallet(seeded_random: random.Random) -> None:
+    dummy_records = DummyDerivationRecords(seeded_random=seeded_random)
+    for i in range(5):
+        dummy_records.generate(i, i * 5)
+    async with DBConnection(1) as wrapper:
+        db = await WalletPuzzleStore.create(wrapper)
+        # Add the records per wallet and verify them
+        for wallet_id, records in dummy_records.records_per_wallet.items():
+            await db.add_derivation_paths(records)
+            for record in records:
+                assert await db.get_derivation_record(record.index, record.wallet_id, record.hardened) == record
+                assert await db.get_wallet_identifier_for_puzzle_hash(record.puzzle_hash) == WalletIdentifier(
+                    record.wallet_id, record.wallet_type
+                )
+        # Remove one wallet after the other and verify before and after each
+        for wallet_id, records in dummy_records.records_per_wallet.items():
+            # Assert the existence again here to make sure the previous removals did not affect other wallet_ids
+            for record in records:
+                assert await db.get_derivation_record(record.index, record.wallet_id, record.hardened) == record
+                assert await db.get_wallet_identifier_for_puzzle_hash(record.puzzle_hash) == WalletIdentifier(
+                    record.wallet_id, record.wallet_type
+                )
+                assert await db.get_last_derivation_path_for_wallet(wallet_id) is not None
+            # Remove the wallet_id and make sure its removed fully
+            await db.delete_wallet(wallet_id)
+            for record in records:
+                assert await db.get_derivation_record(record.index, record.wallet_id, record.hardened) is None
+                assert await db.get_wallet_identifier_for_puzzle_hash(record.puzzle_hash) is None
+                assert await db.get_last_derivation_path_for_wallet(wallet_id) is None
+        assert await db.get_last_derivation_path() is None
+        assert db.last_derivation_index is None
+        assert len(db.last_wallet_derivation_index) == 0