chia-blockchain 2.5.4rc1__py3-none-any.whl → 2.5.5__py3-none-any.whl

chia/_tests/wallet/cat_wallet/test_cat_wallet.py

@@ -2,8 +2,10 @@ from __future__ import annotations
 
 import tempfile
 from pathlib import Path
+from typing import Any, TypeVar
 
 import pytest
+from chia_rs import CoinState
 from chia_rs.sized_bytes import bytes32
 from chia_rs.sized_ints import uint32, uint64
 
@@ -15,27 +17,34 @@ from chia._tests.environments.wallet import (
     WalletTestFramework,
 )
 from chia._tests.util.time_out_assert import time_out_assert, time_out_assert_not_none
-from chia.protocols.wallet_protocol import CoinState
-from chia.rpc.wallet_request_types import GetTransactionMemo, PushTX
 from chia.simulator.simulator_protocol import ReorgProtocol
 from chia.types.blockchain_format.coin import Coin, coin_as_list
-from chia.types.blockchain_format.program import Program
+from chia.types.blockchain_format.program import NIL, Program
 from chia.types.coin_spend import make_spend
 from chia.util.bech32m import encode_puzzle_hash
 from chia.util.db_wrapper import DBWrapper2
 from chia.wallet.cat_wallet.cat_constants import DEFAULT_CATS
 from chia.wallet.cat_wallet.cat_info import LegacyCATInfo
-from chia.wallet.cat_wallet.cat_utils import CAT_MOD, construct_cat_puzzle
+from chia.wallet.cat_wallet.cat_utils import (
+    CAT_MOD,
+    SpendableCAT,
+    construct_cat_puzzle,
+    unsigned_spend_bundle_for_spendable_cats,
+)
 from chia.wallet.cat_wallet.cat_wallet import CATWallet
+from chia.wallet.cat_wallet.r_cat_wallet import RCATWallet
+from chia.wallet.conditions import CreateCoin, UnknownCondition
 from chia.wallet.derivation_record import DerivationRecord
 from chia.wallet.derive_keys import master_pk_to_wallet_pk_unhardened
 from chia.wallet.lineage_proof import LineageProof
 from chia.wallet.puzzles.p2_delegated_puzzle_or_hidden_puzzle import puzzle_hash_for_pk
 from chia.wallet.util.tx_config import DEFAULT_TX_CONFIG
 from chia.wallet.util.wallet_types import WalletType
+from chia.wallet.vc_wallet.vc_drivers import create_revocation_layer
 from chia.wallet.wallet_info import WalletInfo
 from chia.wallet.wallet_interested_store import WalletInterestedStore
 from chia.wallet.wallet_node import WalletNode
+from chia.wallet.wallet_request_types import GetTransactionMemo, PushTX
 from chia.wallet.wallet_state_manager import WalletStateManager
 
 
@@ -43,149 +52,220 @@ def check_wallets(node: WalletNode) -> int:
     return len(node.wallet_state_manager.wallets.keys())
 
 
-@pytest.mark.parametrize(
-    "wallet_environments",
-    [
-        {
-            "num_environments": 1,
-            "blocks_needed": [1],
-        }
-    ],
-    indirect=True,
-)
-@pytest.mark.limit_consensus_modes([ConsensusMode.PLAIN], reason="irrelevant")
-@pytest.mark.anyio
-async def test_cat_creation(wallet_environments: WalletTestFramework) -> None:
-    full_node_api = wallet_environments.full_node
-    wsm = wallet_environments.environments[0].wallet_state_manager
-    wallet = wallet_environments.environments[0].xch_wallet
-    wallet_node = wallet_environments.environments[0].node
-    wallet_environments.environments[0].wallet_aliases = {
-        "xch": 1,
-        "cat": 2,
-    }
-    test_amount = 100
-    test_fee = 10
-    async with wallet.wallet_state_manager.new_action_scope(wallet_environments.tx_config, push=True) as action_scope:
-        cat_wallet = await CATWallet.create_new_cat_wallet(
-            wsm,
-            wallet,
-            {"identifier": "genesis_by_id"},
-            uint64(test_amount),
-            action_scope,
-            fee=uint64(test_fee),
+_T_CATWallet = TypeVar("_T_CATWallet", bound=CATWallet)
+
+
+async def mint_cat(
+    wallet_environments: WalletTestFramework,
+    environment: WalletEnvironment,
+    xch_alias: str,
+    cat_alias: str,
+    amount: uint64,
+    wallet_type: type[_T_CATWallet],
+    tail_nonce: str,
+) -> _T_CATWallet:
+    # (f (q . (() . tail_nonce)))
+    tail = Program.to([5, (1, (None, tail_nonce))])
+    tail_hash = tail.get_tree_hash()
+    async with environment.wallet_state_manager.new_action_scope(
+        wallet_environments.tx_config, push=True
+    ) as action_scope:
+        inner_puzzle_hash = await action_scope.get_puzzle_hash(environment.wallet_state_manager)
+        if wallet_type is RCATWallet:
+            wrapped_inner_puzzle_hash = create_revocation_layer(bytes32.zeros, inner_puzzle_hash).get_tree_hash()
+            extra_args: Any = (bytes32.zeros,)
+        else:
+            wrapped_inner_puzzle_hash = inner_puzzle_hash
+            extra_args = tuple()
+        eve_inner_puzzle = Program.to(
+            (
+                1,
+                [
+                    CreateCoin(wrapped_inner_puzzle_hash, amount, memos=[inner_puzzle_hash]).to_program(),
+                    UnknownCondition(opcode=Program.to(51), args=[NIL, Program.to(-113), tail, NIL]).to_program(),
+                ],
+            )
         )
-        # The next 2 lines are basically a noop, it just adds test coverage
-        cat_wallet = await CATWallet.create(wsm, wallet, cat_wallet.wallet_info)
-        await wsm.add_new_wallet(cat_wallet)
+        eve_cat_puzzle = construct_cat_puzzle(
+            CAT_MOD,
+            tail_hash,
+            eve_inner_puzzle,
+        )
+        eve_cat_puzzle_hash = eve_cat_puzzle.get_tree_hash()
+        await environment.xch_wallet.generate_signed_transaction(
+            amounts=[amount],
+            puzzle_hashes=[eve_cat_puzzle_hash],
+            action_scope=action_scope,
+        )
+        async with action_scope.use() as interface:
+            cat_addition = next(
+                addition
+                for tx in interface.side_effects.transactions
+                for addition in tx.additions
+                if addition.puzzle_hash == eve_cat_puzzle_hash
+            )
+            interface.side_effects.extra_spends.append(
+                unsigned_spend_bundle_for_spendable_cats(
+                    CAT_MOD,
+                    [
+                        SpendableCAT(
+                            cat_addition,
+                            tail_hash,
+                            eve_inner_puzzle,
+                            NIL,
+                        )
+                    ],
+                )
+            )
+
+    cat_wallet = await wallet_type.get_or_create_wallet_for_cat(
+        environment.wallet_state_manager, environment.xch_wallet, tail_hash.hex(), *extra_args
+    )
 
     await wallet_environments.process_pending_states(
         [
-            WalletStateTransition(
+            WalletStateTransition()
+            if env != environment
+            else WalletStateTransition(
                 pre_block_balance_updates={
-                    "xch": {
+                    xch_alias: {
                         "confirmed_wallet_balance": 0,
-                        "unconfirmed_wallet_balance": -test_amount + -test_fee,
-                        "<=#spendable_balance": -test_amount + -test_fee,
-                        "<=#max_send_amount": -test_amount + -test_fee,
+                        "unconfirmed_wallet_balance": -amount,
+                        "<=#spendable_balance": -amount,
+                        "<=#max_send_amount": -amount,
                         ">=#pending_change": 1,  # any amount increase
                         "pending_coin_removal_count": 1,
                     },
-                    "cat": {
+                    cat_alias: {
                         "init": True,
-                        "confirmed_wallet_balance": 0,
-                        "unconfirmed_wallet_balance": test_amount,
-                        "spendable_balance": 0,
-                        "max_send_amount": 0,
-                        "pending_change": test_amount,
-                        "pending_coin_removal_count": 1,
                     },
                 },
                 post_block_balance_updates={
-                    "xch": {
-                        "confirmed_wallet_balance": -test_amount + -test_fee,
+                    xch_alias: {
+                        "confirmed_wallet_balance": -amount,
                         "unconfirmed_wallet_balance": 0,
                         ">=#spendable_balance": 0,
                         ">=#max_send_amount": 0,
                         "<=#pending_change": 1,  # any amount decrease
                         "pending_coin_removal_count": -1,
                     },
-                    "cat": {
-                        "confirmed_wallet_balance": test_amount,
-                        "unconfirmed_wallet_balance": 0,
-                        "spendable_balance": test_amount,
-                        "max_send_amount": test_amount,
-                        "pending_change": -test_amount,
-                        "pending_coin_removal_count": -1,
+                    cat_alias: {
+                        "confirmed_wallet_balance": amount,
+                        "unconfirmed_wallet_balance": amount,
+                        "spendable_balance": amount,
+                        "max_send_amount": amount,
                         "unspent_coin_count": 1,
                     },
                 },
             )
+            for env in wallet_environments.environments
         ]
     )
 
-    # Test migration
-    all_lineage = await cat_wallet.lineage_store.get_all_lineage_proofs()
-    current_info = cat_wallet.wallet_info
-    data_str = bytes(
-        LegacyCATInfo(
-            cat_wallet.cat_info.limitations_program_hash, cat_wallet.cat_info.my_tail, list(all_lineage.items())
-        )
-    ).hex()
-    wallet_info = WalletInfo(current_info.id, current_info.name, current_info.type, data_str)
-    new_cat_wallet = await CATWallet.create(wsm, wallet, wallet_info)
-    assert new_cat_wallet.cat_info.limitations_program_hash == cat_wallet.cat_info.limitations_program_hash
-    assert new_cat_wallet.cat_info.my_tail == cat_wallet.cat_info.my_tail
-    assert await cat_wallet.lineage_store.get_all_lineage_proofs() == all_lineage
+    return cat_wallet
 
-    height = full_node_api.full_node.blockchain.get_peak_height()
-    assert height is not None
-    await full_node_api.reorg_from_index_to_new_index(
-        ReorgProtocol(uint32(height - 1), uint32(height + 1), bytes32(32 * b"1"), None)
+
+@pytest.mark.parametrize(
+    "wallet_environments",
+    [
+        {
+            "num_environments": 1,
+            "blocks_needed": [1],
+        }
+    ],
+    indirect=True,
+)
+@pytest.mark.limit_consensus_modes([ConsensusMode.PLAIN], reason="irrelevant")
+@pytest.mark.parametrize("wallet_type", [CATWallet, RCATWallet])
+@pytest.mark.anyio
+async def test_cat_creation(wallet_environments: WalletTestFramework, wallet_type: type[CATWallet]) -> None:
+    full_node_api = wallet_environments.full_node
+    wsm = wallet_environments.environments[0].wallet_state_manager
+    wallet = wallet_environments.environments[0].xch_wallet
+    wallet_node = wallet_environments.environments[0].node
+    wallet_environments.environments[0].wallet_aliases = {
+        "xch": 1,
+        "cat": 2,
+    }
+    test_amount = uint64(100)
+
+    cat_wallet = await mint_cat(
+        wallet_environments,
+        wallet_environments.environments[0],
+        "xch",
+        "cat",
+        test_amount,
+        wallet_type,
+        "cat wallet",
     )
-    await full_node_api.wait_for_wallet_synced(wallet_node=wallet_node, peak_height=uint32(height + 1))
-    # The "set_remainder" sections here are due to a peculiarity with how the creation method creates an incoming TX
-    # The creation method is for testing purposes only so we're not going to bother fixing it for any real reason
-    await wallet_environments.process_pending_states(
-        [
-            WalletStateTransition(
-                pre_block_balance_updates={
-                    "xch": {
-                        "confirmed_wallet_balance": test_amount + test_fee,
-                        "unconfirmed_wallet_balance": 0,
-                        "<=#spendable_balance": 1,
-                        "<=#max_send_amount": 1,
-                        ">=#pending_change": 1,  # any amount increase
-                        "pending_coin_removal_count": 1,
-                    },
-                    "cat": {
-                        "confirmed_wallet_balance": -test_amount,
-                        "spendable_balance": -test_amount,
-                        "max_send_amount": -test_amount,
-                        "unspent_coin_count": -1,
-                        "set_remainder": True,
-                    },
-                },
-                post_block_balance_updates={
-                    "xch": {
-                        "confirmed_wallet_balance": -test_amount + -test_fee,
-                        "unconfirmed_wallet_balance": 0,
-                        ">=#spendable_balance": 0,
-                        ">=#max_send_amount": 0,
-                        "<=#pending_change": 1,  # any amount decrease
-                        "pending_coin_removal_count": -1,
+
+    # The next 2 lines are basically a noop, it just adds test coverage
+    cat_wallet = await wallet_type.create(wsm, wallet, cat_wallet.wallet_info)
+    await wsm.add_new_wallet(cat_wallet)
+
+    if wallet_type is CATWallet:
+        # Test migration
+        all_lineage = await cat_wallet.lineage_store.get_all_lineage_proofs()
+        current_info = cat_wallet.wallet_info
+        data_str = bytes(
+            LegacyCATInfo(
+                cat_wallet.cat_info.limitations_program_hash, cat_wallet.cat_info.my_tail, list(all_lineage.items())
+            )
+        ).hex()
+        wallet_info = WalletInfo(current_info.id, current_info.name, current_info.type, data_str)
+        new_cat_wallet = await wallet_type.create(wsm, wallet, wallet_info)
+        assert new_cat_wallet.cat_info.limitations_program_hash == cat_wallet.cat_info.limitations_program_hash
+        assert new_cat_wallet.cat_info.my_tail == cat_wallet.cat_info.my_tail
+        assert await cat_wallet.lineage_store.get_all_lineage_proofs() == all_lineage
+
+        height = full_node_api.full_node.blockchain.get_peak_height()
+        assert height is not None
+        await full_node_api.reorg_from_index_to_new_index(
+            ReorgProtocol(uint32(height - 1), uint32(height + 1), bytes32(32 * b"1"), None)
+        )
+        await full_node_api.wait_for_wallet_synced(wallet_node=wallet_node, peak_height=uint32(height + 1))
+        # The "set_remainder" sections here are due to a peculiarity with how the creation method creates an incoming TX
+        # The creation method is for testing purposes only so we're not going to bother fixing it for any real reason
+        await wallet_environments.process_pending_states(
+            [
+                WalletStateTransition(
+                    pre_block_balance_updates={
+                        "xch": {
+                            "confirmed_wallet_balance": test_amount,
+                            "unconfirmed_wallet_balance": 0,
+                            "<=#spendable_balance": 1,
+                            "<=#max_send_amount": 1,
+                            ">=#pending_change": 1,  # any amount increase
+                            "pending_coin_removal_count": 1,
+                        },
+                        "cat": {
+                            "confirmed_wallet_balance": -test_amount,
+                            "spendable_balance": -test_amount,
+                            "max_send_amount": -test_amount,
+                            "unspent_coin_count": -1,
+                            "set_remainder": True,
+                        },
                     },
-                    "cat": {
-                        "confirmed_wallet_balance": test_amount,
-                        "spendable_balance": test_amount,
-                        "max_send_amount": test_amount,
-                        "unspent_coin_count": 1,
-                        "set_remainder": True,
+                    post_block_balance_updates={
+                        "xch": {
+                            "confirmed_wallet_balance": -test_amount,
+                            "unconfirmed_wallet_balance": 0,
+                            ">=#spendable_balance": 0,
+                            ">=#max_send_amount": 0,
+                            "<=#pending_change": 1,  # any amount decrease
+                            "pending_coin_removal_count": -1,
+                        },
+                        "cat": {
+                            "confirmed_wallet_balance": test_amount,
+                            "spendable_balance": test_amount,
+                            "max_send_amount": test_amount,
+                            "unspent_coin_count": 1,
+                            "set_remainder": True,
+                        },
                     },
-                },
-            ),
-        ]
-    )
+                ),
+            ]
+        )
 
 
 @pytest.mark.parametrize(
@@ -201,27 +281,35 @@ async def test_cat_creation(wallet_environments: WalletTestFramework) -> None:
     indirect=True,
 )
 @pytest.mark.limit_consensus_modes([ConsensusMode.PLAIN], reason="irrelevant")
+@pytest.mark.parametrize("wallet_type", [CATWallet, RCATWallet])
 @pytest.mark.anyio
-async def test_cat_creation_unique_lineage_store(wallet_environments: WalletTestFramework) -> None:
-    wsm = wallet_environments.environments[0].wallet_state_manager
-    wallet = wallet_environments.environments[0].xch_wallet
+async def test_cat_creation_unique_lineage_store(
+    wallet_environments: WalletTestFramework, wallet_type: type[CATWallet]
+) -> None:
+    wallet_environments.environments[0].wallet_aliases = {
+        "xch": 1,
+        "cat1": 2,
+        "cat2": 3,
+    }
 
-    async with wsm.new_action_scope(wallet_environments.tx_config, push=True) as action_scope:
-        cat_wallet_1 = await CATWallet.create_new_cat_wallet(
-            wsm,
-            wallet,
-            {"identifier": "genesis_by_id"},
-            uint64(100),
-            action_scope,
-        )
-    async with wsm.new_action_scope(wallet_environments.tx_config, push=True) as action_scope:
-        cat_wallet_2 = await CATWallet.create_new_cat_wallet(
-            wsm,
-            wallet,
-            {"identifier": "genesis_by_id"},
-            uint64(200),
-            action_scope,
-        )
+    cat_wallet_1 = await mint_cat(
+        wallet_environments,
+        wallet_environments.environments[0],
+        "xch",
+        "cat1",
+        uint64(100),
+        wallet_type,
+        "cat wallet 1",
+    )
+    cat_wallet_2 = await mint_cat(
+        wallet_environments,
+        wallet_environments.environments[0],
+        "xch",
+        "cat2",
+        uint64(200),
+        wallet_type,
+        "cat wallet 2",
+    )
 
     proofs_1 = await cat_wallet_1.lineage_store.get_all_lineage_proofs()
     proofs_2 = await cat_wallet_2.lineage_store.get_all_lineage_proofs()
@@ -241,14 +329,14 @@ async def test_cat_creation_unique_lineage_store(wallet_environments: WalletTest
     ],
     indirect=True,
 )
+@pytest.mark.parametrize("wallet_type", [CATWallet, RCATWallet])
 @pytest.mark.anyio
-async def test_cat_spend(wallet_environments: WalletTestFramework) -> None:
+async def test_cat_spend(wallet_environments: WalletTestFramework, wallet_type: type[CATWallet]) -> None:
     # Setup
     env_1: WalletEnvironment = wallet_environments.environments[0]
     env_2: WalletEnvironment = wallet_environments.environments[1]
     wallet_node = env_1.node
     wallet_node_2 = env_2.node
-    wallet = env_1.xch_wallet
     wallet2 = env_2.xch_wallet
     full_node_api = wallet_environments.full_node
 
@@ -261,69 +349,25 @@ async def test_cat_spend(wallet_environments: WalletTestFramework) -> None:
         "cat": 2,
     }
 
-    async with wallet.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
-        cat_wallet = await CATWallet.create_new_cat_wallet(
-            wallet_node.wallet_state_manager,
-            wallet,
-            {"identifier": "genesis_by_id"},
-            uint64(100),
-            action_scope,
-        )
-
-    await wallet_environments.process_pending_states(
-        [
-            WalletStateTransition(
-                pre_block_balance_updates={
-                    "xch": {
-                        "unconfirmed_wallet_balance": -100,
-                        "<=#spendable_balance": -100,
-                        "<=#max_send_amount": -100,
-                        ">=#pending_change": 1,  # any amount increase
-                        "unspent_coin_count": 0,
-                        "pending_coin_removal_count": 1,
-                    },
-                    "cat": {
-                        "init": True,
-                        "confirmed_wallet_balance": 0,
-                        "unconfirmed_wallet_balance": 100,
-                        "spendable_balance": 0,
-                        "pending_change": 100,  # A little weird but technically correct
-                        "max_send_amount": 0,
-                        "unspent_coin_count": 0,
-                        "pending_coin_removal_count": 1,  # The ephemeral eve spend
-                    },
-                },
-                post_block_balance_updates={
-                    "xch": {
-                        "confirmed_wallet_balance": -100,
-                        ">=#spendable_balance": 1,  # any amount increase
-                        ">=#max_send_amount": 1,  # any amount increase
-                        "<=#pending_change": -1,  # any amount decrease
-                        "unspent_coin_count": 0,
-                        "pending_coin_removal_count": -1,
-                    },
-                    "cat": {
-                        "confirmed_wallet_balance": 100,
-                        "spendable_balance": 100,
-                        "pending_change": -100,
-                        "max_send_amount": 100,
-                        "unspent_coin_count": 1,
-                        "pending_coin_removal_count": -1,
-                    },
-                },
-            ),
-            WalletStateTransition(),
-        ]
-    )
+    cat_wallet = await mint_cat(wallet_environments, env_1, "xch", "cat", uint64(100), wallet_type, "cat wallet")
 
     assert cat_wallet.cat_info.limitations_program_hash is not None
     asset_id = cat_wallet.get_asset_id()
 
-    cat_wallet_2 = await CATWallet.get_or_create_wallet_for_cat(wallet_node_2.wallet_state_manager, wallet2, asset_id)
+    if wallet_type is RCATWallet:
+        extra_args: Any = (bytes32.zeros,)
+    else:
+        extra_args = tuple()
+    cat_wallet_2 = await wallet_type.get_or_create_wallet_for_cat(
+        wallet_node_2.wallet_state_manager, wallet2, asset_id, *extra_args
+    )
 
     assert cat_wallet.cat_info.limitations_program_hash == cat_wallet_2.cat_info.limitations_program_hash
 
-    cat_2_hash = await cat_wallet_2.standard_wallet.get_puzzle_hash(new=False)
+    async with cat_wallet_2.wallet_state_manager.new_action_scope(
+        wallet_environments.tx_config, push=True
+    ) as action_scope:
+        cat_2_hash = await action_scope.get_puzzle_hash(cat_wallet_2.wallet_state_manager)
     async with cat_wallet.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
         await cat_wallet.generate_signed_transaction([uint64(60)], [cat_2_hash], action_scope, fee=uint64(1))
     tx_id = None
@@ -335,7 +379,7 @@ async def test_cat_spend(wallet_environments: WalletTestFramework) -> None:
     assert tx_id is not None
     memos = await env_1.rpc_client.get_transaction_memo(GetTransactionMemo(transaction_id=tx_id))
     assert len(memos.coins_with_memos) == 2
-    assert memos.coins_with_memos[1].memos[0] == cat_2_hash
+    assert cat_2_hash in {coin_w_memos.memos[0] for coin_w_memos in memos.coins_with_memos}
 
     await wallet_environments.process_pending_states(
         [
@@ -412,8 +456,11 @@ async def test_cat_spend(wallet_environments: WalletTestFramework) -> None:
     tx_id = coin.name()
     memos = await env_2.rpc_client.get_transaction_memo(GetTransactionMemo(transaction_id=tx_id))
     assert len(memos.coins_with_memos) == 2
-    assert memos.coins_with_memos[1].memos[0] == cat_2_hash
-    cat_hash = await cat_wallet.standard_wallet.get_puzzle_hash(new=False)
+    assert cat_2_hash in {coin_w_memos.memos[0] for coin_w_memos in memos.coins_with_memos}
+    async with cat_wallet.wallet_state_manager.new_action_scope(
+        wallet_environments.tx_config, push=True
+    ) as action_scope:
+        cat_hash = await action_scope.get_puzzle_hash(cat_wallet.wallet_state_manager)
     async with cat_wallet_2.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
         await cat_wallet_2.generate_signed_transaction([uint64(15)], [cat_hash], action_scope)
 
@@ -493,8 +540,9 @@ async def test_cat_spend(wallet_environments: WalletTestFramework) -> None:
     ],
     indirect=True,
 )
+@pytest.mark.parametrize("wallet_type", [CATWallet, RCATWallet])
 @pytest.mark.anyio
-async def test_get_wallet_for_asset_id(wallet_environments: WalletTestFramework) -> None:
+async def test_get_wallet_for_asset_id(wallet_environments: WalletTestFramework, wallet_type: type[CATWallet]) -> None:
     wsm = wallet_environments.environments[0].wallet_state_manager
     wallet = wallet_environments.environments[0].xch_wallet
 
@@ -503,14 +551,9 @@ async def test_get_wallet_for_asset_id(wallet_environments: WalletTestFramework)
         "cat": 2,
     }
 
-    async with wsm.new_action_scope(wallet_environments.tx_config, push=True) as action_scope:
-        cat_wallet = await CATWallet.create_new_cat_wallet(
-            wsm,
-            wallet,
-            {"identifier": "genesis_by_id"},
-            uint64(100),
-            action_scope,
-        )
+    cat_wallet = await mint_cat(
+        wallet_environments, wallet_environments.environments[0], "xch", "cat", uint64(100), wallet_type, "cat wallet"
+    )
 
     await wallet_environments.process_pending_states(
         [
@@ -528,14 +571,16 @@ async def test_get_wallet_for_asset_id(wallet_environments: WalletTestFramework)
     )
 
     asset_id = cat_wallet.get_asset_id()
-    assert cat_wallet.cat_info.my_tail is not None
-    await cat_wallet.set_tail_program(bytes(cat_wallet.cat_info.my_tail).hex())
     assert await wsm.get_wallet_for_asset_id(asset_id) == cat_wallet
 
     # Test that the a default CAT will initialize correctly
     asset = DEFAULT_CATS[next(iter(DEFAULT_CATS))]
     asset_id = asset["asset_id"]
-    cat_wallet_2 = await CATWallet.get_or_create_wallet_for_cat(wsm, wallet, asset_id)
+    if wallet_type is RCATWallet:
+        extra_args: Any = (bytes32.zeros,)
+    else:
+        extra_args = tuple()
+    cat_wallet_2 = await wallet_type.get_or_create_wallet_for_cat(wsm, wallet, asset_id, *extra_args)
     assert cat_wallet_2.get_name() == asset["name"]
     await cat_wallet_2.set_name("Test Name")
     assert cat_wallet_2.get_name() == "Test Name"
@@ -553,12 +598,12 @@ async def test_get_wallet_for_asset_id(wallet_environments: WalletTestFramework)
     ],
     indirect=True,
 )
+@pytest.mark.parametrize("wallet_type", [CATWallet, RCATWallet])
 @pytest.mark.anyio
-async def test_cat_doesnt_see_eve(wallet_environments: WalletTestFramework) -> None:
+async def test_cat_doesnt_see_eve(wallet_environments: WalletTestFramework, wallet_type: type[CATWallet]) -> None:
     # Setup
     env_1: WalletEnvironment = wallet_environments.environments[0]
     env_2: WalletEnvironment = wallet_environments.environments[1]
-    wallet_node = env_1.node
     wallet_node_2 = env_2.node
     wallet = env_1.xch_wallet
     wallet2 = env_2.xch_wallet
@@ -572,50 +617,25 @@ async def test_cat_doesnt_see_eve(wallet_environments: WalletTestFramework) -> N
         "cat": 2,
     }
 
-    async with wallet.wallet_state_manager.new_action_scope(wallet_environments.tx_config, push=True) as action_scope:
-        cat_wallet = await CATWallet.create_new_cat_wallet(
-            wallet_node.wallet_state_manager,
-            wallet,
-            {"identifier": "genesis_by_id"},
-            uint64(100),
-            action_scope,
-        )
-
-    await wallet_environments.process_pending_states(
-        [
-            WalletStateTransition(
-                pre_block_balance_updates={
-                    "xch": {"set_remainder": True},
-                    "cat": {"init": True, "set_remainder": True},
-                },
-                post_block_balance_updates={
-                    "xch": {"set_remainder": True},
-                    "cat": {
-                        "confirmed_wallet_balance": 100,
-                        "unconfirmed_wallet_balance": 0,
-                        "spendable_balance": 100,
-                        "max_send_amount": 100,
-                        "pending_change": -100,
-                        "pending_coin_removal_count": -1,
-                        "unspent_coin_count": 1,
-                    },
-                },
-            ),
-            WalletStateTransition(
-                pre_block_balance_updates={},
-                post_block_balance_updates={},
-            ),
-        ]
-    )
+    cat_wallet = await mint_cat(wallet_environments, env_1, "xch", "cat", uint64(100), wallet_type, "cat wallet")
 
     assert cat_wallet.cat_info.limitations_program_hash is not None
     asset_id = cat_wallet.get_asset_id()
 
-    cat_wallet_2 = await CATWallet.get_or_create_wallet_for_cat(wallet_node_2.wallet_state_manager, wallet2, asset_id)
+    if wallet_type is RCATWallet:
+        extra_args: Any = (bytes32.zeros,)
+    else:
+        extra_args = tuple()
+    cat_wallet_2 = await wallet_type.get_or_create_wallet_for_cat(
+        wallet_node_2.wallet_state_manager, wallet2, asset_id, *extra_args
+    )
 
     assert cat_wallet.cat_info.limitations_program_hash == cat_wallet_2.cat_info.limitations_program_hash
 
-    cat_2_hash = await cat_wallet_2.standard_wallet.get_puzzle_hash(new=False)
+    async with cat_wallet_2.wallet_state_manager.new_action_scope(
+        wallet_environments.tx_config, push=True
+    ) as action_scope:
+        cat_2_hash = await action_scope.get_puzzle_hash(cat_wallet_2.wallet_state_manager)
     async with cat_wallet.wallet_state_manager.new_action_scope(
         wallet_environments.tx_config, push=True
     ) as action_scope:
@@ -746,16 +766,15 @@ async def test_cat_doesnt_see_eve(wallet_environments: WalletTestFramework) -> N
     ],
     indirect=True,
 )
+@pytest.mark.parametrize("wallet_type", [CATWallet, RCATWallet])
 @pytest.mark.anyio
-async def test_cat_spend_multiple(wallet_environments: WalletTestFramework) -> None:
+async def test_cat_spend_multiple(wallet_environments: WalletTestFramework, wallet_type: type[CATWallet]) -> None:
     # Setup
     env_0: WalletEnvironment = wallet_environments.environments[0]
     env_1: WalletEnvironment = wallet_environments.environments[1]
     env_2: WalletEnvironment = wallet_environments.environments[2]
-    wallet_node_0 = env_0.node
     wallet_node_1 = env_1.node
     wallet_node_2 = env_2.node
-    wallet_0 = env_0.xch_wallet
     wallet_1 = env_1.xch_wallet
     wallet_2 = env_2.xch_wallet
 
@@ -772,59 +791,35 @@ async def test_cat_spend_multiple(wallet_environments: WalletTestFramework) -> N
         "cat": 2,
     }
 
-    async with wallet_0.wallet_state_manager.new_action_scope(wallet_environments.tx_config, push=True) as action_scope:
-        cat_wallet_0 = await CATWallet.create_new_cat_wallet(
-            wallet_node_0.wallet_state_manager,
-            wallet_0,
-            {"identifier": "genesis_by_id"},
-            uint64(100),
-            action_scope,
-        )
-
-    await wallet_environments.process_pending_states(
-        [
-            WalletStateTransition(
-                pre_block_balance_updates={
-                    "xch": {"set_remainder": True},
-                    "cat": {"init": True, "set_remainder": True},
-                },
-                post_block_balance_updates={
-                    "xch": {"set_remainder": True},
-                    "cat": {
-                        "confirmed_wallet_balance": 100,
-                        "unconfirmed_wallet_balance": 0,
-                        "spendable_balance": 100,
-                        "max_send_amount": 100,
-                        "pending_change": -100,
-                        "pending_coin_removal_count": -1,
-                        "unspent_coin_count": 1,
-                    },
-                },
-            ),
-            WalletStateTransition(
-                pre_block_balance_updates={},
-                post_block_balance_updates={},
-            ),
-            WalletStateTransition(
-                pre_block_balance_updates={},
-                post_block_balance_updates={},
-            ),
-        ]
-    )
+    cat_wallet_0 = await mint_cat(wallet_environments, env_0, "xch", "cat", uint64(100), wallet_type, "cat wallet")
 
     assert cat_wallet_0.cat_info.limitations_program_hash is not None
     asset_id = cat_wallet_0.get_asset_id()
 
-    cat_wallet_1 = await CATWallet.get_or_create_wallet_for_cat(wallet_node_1.wallet_state_manager, wallet_1, asset_id)
+    if wallet_type is RCATWallet:
+        extra_args: Any = (bytes32.zeros,)
+    else:
+        extra_args = tuple()
+
+    cat_wallet_1 = await wallet_type.get_or_create_wallet_for_cat(
+        wallet_node_1.wallet_state_manager, wallet_1, asset_id, *extra_args
+    )
 
-    cat_wallet_2 = await CATWallet.get_or_create_wallet_for_cat(wallet_node_2.wallet_state_manager, wallet_2, asset_id)
+    cat_wallet_2 = await wallet_type.get_or_create_wallet_for_cat(
+        wallet_node_2.wallet_state_manager, wallet_2, asset_id, *extra_args
+    )
 
     assert cat_wallet_0.cat_info.limitations_program_hash == cat_wallet_1.cat_info.limitations_program_hash
     assert cat_wallet_0.cat_info.limitations_program_hash == cat_wallet_2.cat_info.limitations_program_hash
 
-    cat_1_hash = await cat_wallet_1.standard_wallet.get_puzzle_hash(new=False)
-    cat_2_hash = await cat_wallet_2.standard_wallet.get_puzzle_hash(new=False)
-
+    async with cat_wallet_1.wallet_state_manager.new_action_scope(
+        wallet_environments.tx_config, push=True
+    ) as action_scope:
+        cat_1_hash = await action_scope.get_puzzle_hash(cat_wallet_1.wallet_state_manager)
+    async with cat_wallet_2.wallet_state_manager.new_action_scope(
+        wallet_environments.tx_config, push=True
+    ) as action_scope:
+        cat_2_hash = await action_scope.get_puzzle_hash(cat_wallet_2.wallet_state_manager)
     async with cat_wallet_0.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
         await cat_wallet_0.generate_signed_transaction([uint64(60), uint64(20)], [cat_1_hash, cat_2_hash], action_scope)
 
@@ -905,8 +900,10 @@ async def test_cat_spend_multiple(wallet_environments: WalletTestFramework) -> N
         ]
     )
 
-    cat_hash = await cat_wallet_0.standard_wallet.get_puzzle_hash(new=False)
-
+    async with cat_wallet_0.wallet_state_manager.new_action_scope(
+        wallet_environments.tx_config, push=True
+    ) as action_scope:
+        cat_hash = await action_scope.get_puzzle_hash(cat_wallet_0.wallet_state_manager)
     async with cat_wallet_1.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
         await cat_wallet_1.generate_signed_transaction([uint64(15)], [cat_hash], action_scope)
 
@@ -1066,63 +1063,32 @@ async def test_cat_spend_multiple(wallet_environments: WalletTestFramework) -> N
     ],
     indirect=True,
 )
+@pytest.mark.parametrize("wallet_type", [CATWallet, RCATWallet])
 @pytest.mark.anyio
-async def test_cat_max_amount_send(wallet_environments: WalletTestFramework) -> None:
+async def test_cat_max_amount_send(wallet_environments: WalletTestFramework, wallet_type: type[CATWallet]) -> None:
     # Setup
     env: WalletEnvironment = wallet_environments.environments[0]
-    wallet_node = env.node
-    wallet = env.xch_wallet
 
     env.wallet_aliases = {
         "xch": 1,
         "cat": 2,
     }
 
-    async with wallet.wallet_state_manager.new_action_scope(wallet_environments.tx_config, push=True) as action_scope:
-        cat_wallet = await CATWallet.create_new_cat_wallet(
-            wallet_node.wallet_state_manager,
-            wallet,
-            {"identifier": "genesis_by_id"},
-            uint64(100000),
-            action_scope,
-        )
-
-    await wallet_environments.process_pending_states(
-        [
-            WalletStateTransition(
-                pre_block_balance_updates={
-                    "xch": {"set_remainder": True},
-                    "cat": {"init": True, "set_remainder": True},
-                },
-                post_block_balance_updates={
-                    "xch": {"set_remainder": True},
-                    "cat": {
-                        "confirmed_wallet_balance": 100000,
-                        "unconfirmed_wallet_balance": 0,
-                        "spendable_balance": 100000,
-                        "max_send_amount": 100000,
-                        "pending_change": -100000,
-                        "pending_coin_removal_count": -1,
-                        "unspent_coin_count": 1,
-                    },
-                },
-            )
-        ]
-    )
+    cat_wallet = await mint_cat(wallet_environments, env, "xch", "cat", uint64(100000), wallet_type, "cat wallet")
 
     assert cat_wallet.cat_info.limitations_program_hash is not None
 
-    cat_2 = await cat_wallet.standard_wallet.get_puzzle(new=False)
-    cat_2_hash = cat_2.get_tree_hash()
-    amounts = []
-    puzzle_hashes = []
-    for i in range(1, 50):
-        amounts.append(uint64(i))
-        puzzle_hashes.append(cat_2_hash)
-    spent_coin = (await cat_wallet.get_cat_spendable_coins())[0].coin
     async with cat_wallet.wallet_state_manager.new_action_scope(
         wallet_environments.tx_config, push=True
     ) as action_scope:
+        cat_2 = await action_scope.get_puzzle(cat_wallet.wallet_state_manager)
+        cat_2_hash = cat_2.get_tree_hash()
+        amounts = []
+        puzzle_hashes = []
+        for i in range(1, 50):
+            amounts.append(uint64(i))
+            puzzle_hashes.append(cat_2_hash)
+        spent_coin = (await cat_wallet.get_cat_spendable_coins())[0].coin
         await cat_wallet.generate_signed_transaction(amounts, puzzle_hashes, action_scope, coins={spent_coin})
 
     await wallet_environments.process_pending_states(
@@ -1157,7 +1123,13 @@ async def test_cat_max_amount_send(wallet_environments: WalletTestFramework) ->
         spendable_name_set = set()
         for record in spendable:
             spendable_name_set.add(record.coin.name())
-        puzzle_hash = construct_cat_puzzle(CAT_MOD, cat_wallet.cat_info.limitations_program_hash, cat_2).get_tree_hash()
+        if wallet_type is RCATWallet:
+            inner_puzzle = create_revocation_layer(bytes32.zeros, cat_2_hash)
+        else:
+            inner_puzzle = cat_2
+        puzzle_hash = construct_cat_puzzle(
+            CAT_MOD, cat_wallet.cat_info.limitations_program_hash, inner_puzzle
+        ).get_tree_hash()
         for i in range(1, 50):
             coin = Coin(spent_coin.name(), puzzle_hash, uint64(i))
             if coin.name() not in spendable_name_set:
@@ -1206,8 +1178,9 @@ async def test_cat_max_amount_send(wallet_environments: WalletTestFramework) ->
     ],
     indirect=True,
 )
+@pytest.mark.parametrize("wallet_type", [CATWallet, RCATWallet])
 @pytest.mark.anyio
-async def test_cat_hint(wallet_environments: WalletTestFramework) -> None:
+async def test_cat_hint(wallet_environments: WalletTestFramework, wallet_type: type[CATWallet]) -> None:
     # Setup
     env_1: WalletEnvironment = wallet_environments.environments[0]
     env_2: WalletEnvironment = wallet_environments.environments[1]
@@ -1227,45 +1200,12 @@ async def test_cat_hint(wallet_environments: WalletTestFramework) -> None:
 
     autodiscovery = wallet_node_1.config["automatically_add_unknown_cats"]
 
-    async with wallet_1.wallet_state_manager.new_action_scope(wallet_environments.tx_config, push=True) as action_scope:
-        cat_wallet = await CATWallet.create_new_cat_wallet(
-            wallet_node_1.wallet_state_manager,
-            wallet_1,
-            {"identifier": "genesis_by_id"},
-            uint64(100),
-            action_scope,
-        )
-
-    await wallet_environments.process_pending_states(
-        [
-            WalletStateTransition(
-                pre_block_balance_updates={
-                    "xch": {"set_remainder": True},
-                    "cat": {"init": True, "set_remainder": True},
-                },
-                post_block_balance_updates={
-                    "xch": {"set_remainder": True},
-                    "cat": {
-                        "confirmed_wallet_balance": 100,
-                        "unconfirmed_wallet_balance": 0,
-                        "spendable_balance": 100,
-                        "max_send_amount": 100,
-                        "pending_change": -100,
-                        "pending_coin_removal_count": -1,
-                        "unspent_coin_count": 1,
-                    },
-                },
-            ),
-            WalletStateTransition(
-                pre_block_balance_updates={},
-                post_block_balance_updates={},
-            ),
-        ]
-    )
+    cat_wallet = await mint_cat(wallet_environments, env_1, "xch", "cat", uint64(100), wallet_type, "cat wallet")
 
     assert cat_wallet.cat_info.limitations_program_hash is not None
 
-    cat_2_hash = await wallet_2.get_new_puzzlehash()
+    async with wallet_2.wallet_state_manager.new_action_scope(wallet_environments.tx_config, push=True) as action_scope:
+        cat_2_hash = await action_scope.get_puzzle_hash(wallet_2.wallet_state_manager)
     async with wallet_1.wallet_state_manager.new_action_scope(wallet_environments.tx_config, push=True) as action_scope:
         await cat_wallet.generate_signed_transaction([uint64(60)], [cat_2_hash], action_scope, memos=[[cat_2_hash]])
 
@@ -1387,7 +1327,10 @@ async def test_cat_hint(wallet_environments: WalletTestFramework) -> None:
     cat_wallet_2 = wallet_node_2.wallet_state_manager.wallets[uint32(2)]
     assert isinstance(cat_wallet_2, CATWallet)
 
-    cat_hash = await cat_wallet.standard_wallet.get_puzzle_hash(new=False)
+    async with cat_wallet.wallet_state_manager.new_action_scope(
+        wallet_environments.tx_config, push=True
+    ) as action_scope:
+        cat_hash = await action_scope.get_puzzle_hash(cat_wallet.wallet_state_manager)
     async with cat_wallet_2.wallet_state_manager.new_action_scope(
         wallet_environments.tx_config, push=True
     ) as action_scope:
@@ -1449,8 +1392,9 @@ async def test_cat_hint(wallet_environments: WalletTestFramework) -> None:
     ],
     indirect=True,
 )
+@pytest.mark.parametrize("wallet_type", [CATWallet, RCATWallet])
 @pytest.mark.anyio
-async def test_cat_change_detection(wallet_environments: WalletTestFramework) -> None:
+async def test_cat_change_detection(wallet_environments: WalletTestFramework, wallet_type: type[CATWallet]) -> None:
     full_node_api = wallet_environments.full_node
     env = wallet_environments.environments[0]
     wsm = env.wallet_state_manager
@@ -1465,6 +1409,8 @@ async def test_cat_change_detection(wallet_environments: WalletTestFramework) ->
     # We should pick up this coin as balance even though it is unhinted because it is "change"
     pubkey_unhardened = master_pk_to_wallet_pk_unhardened(wsm.root_pubkey, uint32(100000000))
     inner_puzhash = puzzle_hash_for_pk(pubkey_unhardened)
+    if wallet_type is RCATWallet:
+        inner_puzhash = create_revocation_layer(bytes32.zeros, inner_puzhash).get_tree_hash()
     puzzlehash_unhardened = construct_cat_puzzle(
         CAT_MOD, Program.to(None).get_tree_hash(), inner_puzhash
     ).get_tree_hash_precalc(inner_puzhash)
@@ -1473,7 +1419,8 @@ async def test_cat_change_detection(wallet_environments: WalletTestFramework) ->
     )
     # Insert the derivation record before the wallet exists so that it is not subscribed to
     await wsm.puzzle_store.add_derivation_paths([change_derivation])
-    our_puzzle = await wallet.get_new_puzzle()
+    async with wallet.wallet_state_manager.new_action_scope(wallet_environments.tx_config, push=True) as action_scope:
+        our_puzzle = await action_scope.get_puzzle(wallet.wallet_state_manager)
     cat_puzzle = construct_cat_puzzle(
         CAT_MOD,
         Program.to(None).get_tree_hash(),
@@ -1626,7 +1573,6 @@ async def test_cat_melt_balance(wallet_environments: WalletTestFramework) -> Non
     env = wallet_environments.environments[0]
     wallet = env.xch_wallet
     simulator = wallet_environments.full_node
-    wallet_ph = await wallet.get_puzzle_hash(new=False)
 
     env.wallet_aliases = {
         "xch": 1,
@@ -1653,18 +1599,22 @@ async def test_cat_melt_balance(wallet_environments: WalletTestFramework) -> Non
 
     tx_amount = 10
 
-    spend_to_wallet = unsigned_spend_bundle_for_spendable_cats(
-        CAT_MOD,
-        [
-            SpendableCAT(
-                coin=cat_coin,
-                limitations_program_hash=ACS_TAIL_HASH,
-                inner_puzzle=ACS,
-                inner_solution=Program.to([[51, wallet_ph, tx_amount, [wallet_ph]], [51, None, -113, ACS_TAIL, None]]),
-                extra_delta=tx_amount - cat_coin.amount,
-            )
-        ],
-    )
+    async with env.wallet_state_manager.new_action_scope(wallet_environments.tx_config, push=True) as action_scope:
+        wallet_ph = await action_scope.get_puzzle_hash(env.wallet_state_manager)
+        spend_to_wallet = unsigned_spend_bundle_for_spendable_cats(
+            CAT_MOD,
+            [
+                SpendableCAT(
+                    coin=cat_coin,
+                    limitations_program_hash=ACS_TAIL_HASH,
+                    inner_puzzle=ACS,
+                    inner_solution=Program.to(
+                        [[51, wallet_ph, tx_amount, [wallet_ph]], [51, None, -113, ACS_TAIL, None]]
+                    ),
+                    extra_delta=tx_amount - cat_coin.amount,
+                )
+            ],
+        )
     await env.rpc_client.push_tx(PushTX(spend_to_wallet))
     await time_out_assert(10, simulator.tx_id_in_mempool, True, spend_to_wallet.name())
 
@@ -1691,7 +1641,7 @@ async def test_cat_melt_balance(wallet_environments: WalletTestFramework) -> Non
     assert isinstance(cat_wallet, CATWallet)
 
     # Let's test that continuing to melt this CAT results in the correct balance changes
-    for _ in range(0, 5):
+    for _ in range(5):
         tx_amount -= 1
         new_coin = (await cat_wallet.get_cat_spendable_coins())[0].coin
         new_spend = unsigned_spend_bundle_for_spendable_cats(
@@ -1748,54 +1698,26 @@ async def test_cat_melt_balance(wallet_environments: WalletTestFramework) -> Non
     ],
     indirect=True,
 )
+@pytest.mark.parametrize("wallet_type", [CATWallet, RCATWallet])
 @pytest.mark.limit_consensus_modes([ConsensusMode.PLAIN], reason="irrelevant")
 @pytest.mark.anyio
-async def test_cat_puzzle_hashes(wallet_environments: WalletTestFramework) -> None:
+async def test_cat_puzzle_hashes(wallet_environments: WalletTestFramework, wallet_type: type[CATWallet]) -> None:
     env = wallet_environments.environments[0]
-    wallet = env.xch_wallet
 
     env.wallet_aliases = {
         "xch": 1,
         "cat": 2,
     }
 
-    async with env.wallet_state_manager.new_action_scope(wallet_environments.tx_config, push=True) as action_scope:
-        cat_wallet = await CATWallet.create_new_cat_wallet(
-            env.node.wallet_state_manager,
-            wallet,
-            {"identifier": "genesis_by_id"},
-            uint64(100),
-            action_scope,
-        )
-
-    await wallet_environments.process_pending_states(
-        [
-            WalletStateTransition(
-                pre_block_balance_updates={
-                    "xch": {"set_remainder": True},
-                    "cat": {"init": True, "set_remainder": True},
-                },
-                post_block_balance_updates={
-                    "xch": {"set_remainder": True},
-                    "cat": {
-                        "confirmed_wallet_balance": 100,
-                        "unconfirmed_wallet_balance": 0,
-                        "spendable_balance": 100,
-                        "max_send_amount": 100,
-                        "pending_change": -100,
-                        "pending_coin_removal_count": -1,
-                        "unspent_coin_count": 1,
-                    },
-                },
-            ),
-        ]
-    )
+    cat_wallet = await mint_cat(wallet_environments, env, "xch", "cat", uint64(100), wallet_type, "cat wallet")
 
     # Test that we attempt a new puzzle hash here even though everything says we shouldn't
     with pytest.raises(NewPuzzleHashError):
         async with env.wallet_state_manager.new_action_scope(wallet_environments.tx_config, push=True) as action_scope:
             await cat_wallet.generate_signed_transaction(
-                [uint64(50)], [await cat_wallet.standard_wallet.get_puzzle_hash(new=False)], action_scope
+                [uint64(50)],
+                [await action_scope.get_puzzle_hash(cat_wallet.wallet_state_manager)],
+                action_scope,
             )
 
     # Test new puzzle hash getting
@@ -1818,8 +1740,6 @@ async def test_cat_puzzle_hashes(wallet_environments: WalletTestFramework) -> No
         await env.wallet_state_manager.puzzle_store.add_derivation_paths([dr])
 
     # Then we make sure that even though we asked for a used puzzle hash, it still gives us an unused one
-    unused_count = await env.wallet_state_manager.puzzle_store.get_unused_count(uint32(env.wallet_aliases["cat"]))
+    unused_count = await env.wallet_state_manager.puzzle_store.get_used_count(uint32(env.wallet_aliases["cat"]))
     await cat_wallet.get_cat_puzzle_hash(new=False)
-    assert unused_count < await env.wallet_state_manager.puzzle_store.get_unused_count(
-        uint32(env.wallet_aliases["cat"])
-    )
+    assert unused_count < await env.wallet_state_manager.puzzle_store.get_used_count(uint32(env.wallet_aliases["cat"]))