chia-blockchain 2.5.0rc2__py3-none-any.whl → 2.5.1rc2__py3-none-any.whl

chia/_tests/wallet/did_wallet/test_did.py

@@ -2,16 +2,19 @@ from __future__ import annotations
 
 import dataclasses
 import json
-from typing import Optional
 
 import pytest
 from chia_rs import AugSchemeMPL, G1Element, G2Element
 
+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.rpc.wallet_request_types import DIDGetCurrentCoinInfo, DIDGetRecoveryInfo
 from chia.rpc.wallet_rpc_api import WalletRpcApi
+from chia.server.server import ChiaServer
+from chia.simulator.block_tools import BlockTools
+from chia.simulator.full_node_simulator import FullNodeSimulator
 from chia.simulator.simulator_protocol import FarmNewBlockProtocol
 from chia.types.blockchain_format.program import Program
 from chia.types.blockchain_format.sized_bytes import bytes32
@@ -26,6 +29,7 @@ from chia.wallet.singleton import create_singleton_puzzle
 from chia.wallet.util.address_type import AddressType
 from chia.wallet.util.tx_config import DEFAULT_TX_CONFIG
 from chia.wallet.util.wallet_types import WalletType
+from chia.wallet.wallet_node import WalletNode
 from chia.wallet.wallet_spend_bundle import WalletSpendBundle
 
 
@@ -38,14 +42,16 @@ def get_parent_num(did_wallet: DIDWallet):
 
 
 class TestDIDWallet:
+    #  TODO: See Issue CHIA-1544
+    #  This test should be ported to WalletTestFramework once we can replace keys in the wallet node
     @pytest.mark.parametrize(
         "trusted",
         [True, False],
     )
     @pytest.mark.anyio
     async def test_creation_from_coin_spend(
-        self, self_hostname, two_nodes_two_wallets_with_same_keys: OldSimulatorsAndWallets, trusted
-    ):
+        self, self_hostname: str, two_nodes_two_wallets_with_same_keys: OldSimulatorsAndWallets, trusted: bool
+    ) -> None:
         """
         Verify that DIDWallet.create_new_did_wallet_from_coin_spend() is called after Singleton creation on
         the blockchain, and that the wallet is created in the second wallet node.
@@ -110,7 +116,7 @@ class TestDIDWallet:
             == json.loads(all_node_1_wallets[1].data)["current_inner"]
         )
 
-    # TODO: Porting this test to this fixture revealed some balance peculiarities. Fix them.
+    @pytest.mark.limit_consensus_modes(allowed=[ConsensusMode.PLAIN], reason="irrelevant")
     @pytest.mark.parametrize(
         "wallet_environments",
         [
@@ -163,8 +169,8 @@ class TestDIDWallet:
                         "did": {
                             "init": True,
                             "unconfirmed_wallet_balance": 101,
-                            "pending_change": 202,  # TODO: this is not correct, fix this
-                            "pending_coin_removal_count": 2,  # TODO: this might not be correct
+                            "pending_change": 101,
+                            "pending_coin_removal_count": 1,
                         },
                     },
                     post_block_balance_updates={
@@ -180,8 +186,8 @@ class TestDIDWallet:
                             "spendable_balance": 101,
                             "max_send_amount": 101,
                             "unspent_coin_count": 1,
-                            "pending_change": -202,  # TODO: this is not correct, fix this
-                            "pending_coin_removal_count": -2,  # TODO: this might not be correct
+                            "pending_change": -101,
+                            "pending_coin_removal_count": -1,
                         },
                     },
                 ),
@@ -224,8 +230,8 @@ class TestDIDWallet:
                         "did": {
                             "init": True,
                             "unconfirmed_wallet_balance": 201,
-                            "pending_change": 402,  # TODO: this is not correct, fix this
-                            "pending_coin_removal_count": 2,  # TODO: this might not be correct
+                            "pending_change": 201,
+                            "pending_coin_removal_count": 1,
                         },
                     },
                     post_block_balance_updates={
@@ -241,8 +247,8 @@ class TestDIDWallet:
                             "spendable_balance": 201,
                             "max_send_amount": 201,
                             "unspent_coin_count": 1,
-                            "pending_change": -402,  # TODO: this is not correct, fix this
-                            "pending_coin_removal_count": -2,  # TODO: this might not be correct
+                            "pending_change": -201,
+                            "pending_coin_removal_count": -1,
                         },
                     },
                 ),
@@ -288,6 +294,7 @@ class TestDIDWallet:
                             "spendable_balance": -101,
                             "pending_change": 101,
                             "pending_coin_removal_count": 1,
+                            "max_send_amount": -101,
                         }
                     },
                     post_block_balance_updates={
@@ -295,6 +302,7 @@ class TestDIDWallet:
                             "spendable_balance": 101,
                             "pending_change": -101,
                             "pending_coin_removal_count": -1,
+                            "max_send_amount": 101,
                         }
                     },
                 ),
@@ -394,14 +402,13 @@ class TestDIDWallet:
                         "did": {
                             "unconfirmed_wallet_balance": -201,
                             "spendable_balance": -201,
-                            # "max_send_amount": -201,  # TODO: Uncomment this
+                            "max_send_amount": -201,
                             "pending_coin_removal_count": 1,
                         }
                     },
                     post_block_balance_updates={
                         "did": {
                             "confirmed_wallet_balance": -201,
-                            "max_send_amount": -201,  # TODO: Delete this when uncommented above
                             "unspent_coin_count": -1,
                             "pending_coin_removal_count": -1,
                         }
@@ -421,97 +428,178 @@ class TestDIDWallet:
         for wallet in [did_wallet_0, did_wallet_1]:
             assert wallet.wallet_state_manager.wallets[wallet.id()] == wallet
 
-    @pytest.mark.parametrize(
-        "trusted",
-        [True, False],
-    )
+    @pytest.mark.limit_consensus_modes(allowed=[ConsensusMode.PLAIN], reason="irrelevant")
+    @pytest.mark.parametrize("wallet_environments", [{"num_environments": 2, "blocks_needed": [1, 1]}], indirect=True)
     @pytest.mark.anyio
-    async def test_did_recovery_with_multiple_backup_dids(self, self_hostname, two_wallet_nodes, trusted):
-        full_nodes, wallets, _ = two_wallet_nodes
-        full_node_api = full_nodes[0]
-        server_1 = full_node_api.server
-        wallet_node, server_2 = wallets[0]
-        wallet_node_2, server_3 = wallets[1]
-        wallet = wallet_node.wallet_state_manager.main_wallet
-        wallet2 = wallet_node_2.wallet_state_manager.main_wallet
-
-        if trusted:
-            wallet_node.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.config["trusted_peers"] = {}
-            wallet_node_2.config["trusted_peers"] = {}
-        await server_2.start_client(PeerInfo(self_hostname, server_1.get_port()), None)
-        await server_3.start_client(PeerInfo(self_hostname, server_1.get_port()), None)
+    async def test_did_recovery_with_multiple_backup_dids(self, wallet_environments: WalletTestFramework):
+        env_0 = wallet_environments.environments[0]
+        env_1 = wallet_environments.environments[1]
+        wallet_node_0 = env_0.node
+        wallet_node_1 = env_1.node
+        wallet_0 = env_0.xch_wallet
+        wallet_1 = env_1.xch_wallet
 
-        await full_node_api.farm_blocks_to_wallet(1, wallet)
-        await full_node_api.farm_blocks_to_wallet(1, wallet2)
+        env_0.wallet_aliases = {
+            "xch": 1,
+            "did": 2,
+        }
+        env_1.wallet_aliases = {
+            "xch": 1,
+            "did": 2,
+        }
 
-        async with wallet.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
+        async with wallet_0.wallet_state_manager.new_action_scope(
+            wallet_environments.tx_config, push=True
+        ) as action_scope:
             did_wallet: DIDWallet = await DIDWallet.create_new_did_wallet(
-                wallet_node.wallet_state_manager, wallet, uint64(101), action_scope
+                wallet_node_0.wallet_state_manager, wallet_0, uint64(101), action_scope
             )
         assert did_wallet.get_name() == "Profile 1"
-        await full_node_api.process_transaction_records(records=action_scope.side_effects.transactions)
-        await full_node_api.wait_for_wallets_synced(wallet_nodes=[wallet_node, wallet_node_2])
-
-        await time_out_assert(15, did_wallet.get_confirmed_balance, 101)
-        await time_out_assert(15, did_wallet.get_unconfirmed_balance, 101)
-
-        recovery_list = [bytes.fromhex(did_wallet.get_my_DID())]
+        recovery_list = [bytes32.from_hexstr(did_wallet.get_my_DID())]
 
-        async with wallet2.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
+        async with wallet_1.wallet_state_manager.new_action_scope(
+            wallet_environments.tx_config, push=True
+        ) as action_scope:
             did_wallet_2: DIDWallet = await DIDWallet.create_new_did_wallet(
-                wallet_node_2.wallet_state_manager, wallet2, uint64(101), action_scope, recovery_list
+                wallet_node_1.wallet_state_manager, wallet_1, uint64(101), action_scope, recovery_list
             )
 
-        await full_node_api.process_transaction_records(records=action_scope.side_effects.transactions)
-        await full_node_api.wait_for_wallets_synced(wallet_nodes=[wallet_node, wallet_node_2])
-
-        await time_out_assert(15, did_wallet_2.get_confirmed_balance, 101)
-        await time_out_assert(15, did_wallet_2.get_unconfirmed_balance, 101)
-
+        await wallet_environments.process_pending_states(
+            [
+                WalletStateTransition(
+                    pre_block_balance_updates={
+                        "xch": {
+                            "set_remainder": True,
+                        },
+                        "did": {
+                            "init": True,
+                            "unconfirmed_wallet_balance": 101,
+                            "pending_change": 101,
+                            "pending_coin_removal_count": 1,
+                        },
+                    },
+                    post_block_balance_updates={
+                        "xch": {
+                            "set_remainder": True,
+                        },
+                        "did": {
+                            "confirmed_wallet_balance": 101,
+                            "spendable_balance": 101,
+                            "max_send_amount": 101,
+                            "unspent_coin_count": 1,
+                            "pending_change": -101,
+                            "pending_coin_removal_count": -1,
+                        },
+                    },
+                ),
+                WalletStateTransition(
+                    pre_block_balance_updates={
+                        "xch": {
+                            "set_remainder": True,
+                        },
+                        "did": {
+                            "init": True,
+                            "unconfirmed_wallet_balance": 101,
+                            "pending_change": 101,
+                            "pending_coin_removal_count": 1,
+                        },
+                    },
+                    post_block_balance_updates={
+                        "xch": {
+                            "set_remainder": True,
+                        },
+                        "did": {
+                            "confirmed_wallet_balance": 101,
+                            "spendable_balance": 101,
+                            "max_send_amount": 101,
+                            "unspent_coin_count": 1,
+                            "pending_change": -101,
+                            "pending_coin_removal_count": -1,
+                        },
+                    },
+                ),
+            ]
+        )
         assert did_wallet_2.did_info.backup_ids == recovery_list
 
-        recovery_list.append(bytes.fromhex(did_wallet_2.get_my_DID()))
+        recovery_list.append(bytes32.from_hexstr(did_wallet_2.get_my_DID()))
 
-        async with wallet2.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
+        async with wallet_1.wallet_state_manager.new_action_scope(
+            wallet_environments.tx_config, push=True
+        ) as action_scope:
             did_wallet_3: DIDWallet = await DIDWallet.create_new_did_wallet(
-                wallet_node_2.wallet_state_manager, wallet2, uint64(201), action_scope, recovery_list
+                wallet_node_1.wallet_state_manager, wallet_1, uint64(201), action_scope, recovery_list
             )
 
-        await full_node_api.process_transaction_records(records=action_scope.side_effects.transactions)
-        await full_node_api.wait_for_wallets_synced(wallet_nodes=[wallet_node, wallet_node_2])
+        env_1.wallet_aliases["did_2"] = 3
 
-        assert did_wallet_3.did_info.backup_ids == recovery_list
-        await time_out_assert(15, did_wallet_3.get_confirmed_balance, 201)
-        await time_out_assert(15, did_wallet_3.get_unconfirmed_balance, 201)
+        await wallet_environments.process_pending_states(
+            [
+                WalletStateTransition(
+                    pre_block_balance_updates={},
+                    post_block_balance_updates={},
+                ),
+                WalletStateTransition(
+                    pre_block_balance_updates={
+                        "xch": {
+                            "set_remainder": True,
+                        },
+                        "did": {
+                            "set_remainder": True,
+                        },
+                        "did_2": {
+                            "init": True,
+                            "unconfirmed_wallet_balance": 201,
+                            "pending_change": 201,
+                            "pending_coin_removal_count": 1,
+                        },
+                    },
+                    post_block_balance_updates={
+                        "xch": {
+                            "set_remainder": True,
+                        },
+                        "did": {
+                            "set_remainder": True,
+                        },
+                        "did_2": {
+                            "confirmed_wallet_balance": 201,
+                            "spendable_balance": 201,
+                            "max_send_amount": 201,
+                            "unspent_coin_count": 1,
+                            "pending_change": -201,
+                            "pending_coin_removal_count": -1,
+                        },
+                    },
+                ),
+            ]
+        )
         coin = await did_wallet_3.get_coin()
 
         backup_data = did_wallet_3.create_backup()
 
-        async with wallet_node.wallet_state_manager.lock:
+        async with wallet_node_0.wallet_state_manager.lock:
             did_wallet_4 = await DIDWallet.create_new_did_wallet_from_recovery(
-                wallet_node.wallet_state_manager,
-                wallet,
+                wallet_node_0.wallet_state_manager,
+                wallet_0,
                 backup_data,
             )
         assert did_wallet_4.get_name() == "Profile 2"
+        env_0.wallet_aliases["did_2"] = 3
 
         pubkey = (
             await did_wallet_4.wallet_state_manager.get_unused_derivation_record(did_wallet_2.wallet_info.id)
         ).pubkey
         new_ph = did_wallet_4.did_info.temp_puzhash
-        async with did_wallet.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
+        async with did_wallet.wallet_state_manager.new_action_scope(
+            wallet_environments.tx_config, push=True
+        ) as action_scope:
             message_spend_bundle, attest1 = await did_wallet.create_attestment(
                 coin.name(), new_ph, pubkey, action_scope
             )
 
-        async with did_wallet_2.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope_2:
+        async with did_wallet_2.wallet_state_manager.new_action_scope(
+            wallet_environments.tx_config, push=True
+        ) as action_scope_2:
             message_spend_bundle2, attest2 = await did_wallet_2.create_attestment(
                 coin.name(), new_ph, pubkey, action_scope_2
             )
@@ -524,233 +612,486 @@ class TestDIDWallet:
         ) = await did_wallet_4.load_attest_files_for_recovery_spend([attest1, attest2])
         assert message_spend_bundle == test_message_spend_bundle
 
-        await full_node_api.process_transaction_records(
-            records=[*action_scope.side_effects.transactions, *action_scope_2.side_effects.transactions]
+        await wallet_environments.process_pending_states(
+            [
+                WalletStateTransition(
+                    pre_block_balance_updates={
+                        "did": {
+                            "spendable_balance": -101,
+                            "pending_change": 101,
+                            "max_send_amount": -101,
+                            "pending_coin_removal_count": 1,
+                            "set_remainder": True,
+                        },
+                        "did_2": {
+                            "init": True,
+                            "unconfirmed_wallet_balance": 0,
+                            "pending_change": 0,
+                            "pending_coin_removal_count": 0,
+                        },
+                    },
+                    post_block_balance_updates={
+                        "did": {
+                            "spendable_balance": 101,
+                            "pending_change": -101,
+                            "max_send_amount": 101,
+                            "pending_coin_removal_count": -1,
+                        },
+                        "did_2": {
+                            "confirmed_wallet_balance": 0,
+                            "spendable_balance": 0,
+                            "max_send_amount": 0,
+                            "unspent_coin_count": 0,
+                            "pending_change": 0,
+                            "pending_coin_removal_count": 0,
+                        },
+                    },
+                ),
+                WalletStateTransition(
+                    pre_block_balance_updates={
+                        "did": {
+                            "spendable_balance": -101,
+                            "pending_change": 101,
+                            "max_send_amount": -101,
+                            "pending_coin_removal_count": 1,
+                            "set_remainder": True,
+                        },
+                    },
+                    post_block_balance_updates={
+                        "did": {
+                            "spendable_balance": 101,
+                            "pending_change": -101,
+                            "max_send_amount": 101,
+                            "pending_coin_removal_count": -1,
+                        },
+                    },
+                ),
+            ]
         )
-        await full_node_api.wait_for_wallets_synced(wallet_nodes=[wallet_node, wallet_node_2])
-        await time_out_assert(15, did_wallet_4.get_confirmed_balance, 0)
-        await time_out_assert(15, did_wallet_4.get_unconfirmed_balance, 0)
-        async with did_wallet_4.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
+
+        async with did_wallet_4.wallet_state_manager.new_action_scope(
+            wallet_environments.tx_config, push=True
+        ) as action_scope:
             await did_wallet_4.recovery_spend(coin, new_ph, test_info_list, pubkey, message_spend_bundle, action_scope)
-        await full_node_api.process_transaction_records(action_scope.side_effects.transactions)
-        await full_node_api.wait_for_wallets_synced(wallet_nodes=[wallet_node, wallet_node_2])
 
-        await time_out_assert(15, did_wallet_4.get_confirmed_balance, 201)
-        await time_out_assert(15, did_wallet_4.get_unconfirmed_balance, 201)
-        await time_out_assert(15, did_wallet_3.get_confirmed_balance, 0)
-        await time_out_assert(15, did_wallet_3.get_unconfirmed_balance, 0)
+        await wallet_environments.process_pending_states(
+            [
+                WalletStateTransition(
+                    pre_block_balance_updates={
+                        "did_2": {
+                            "unconfirmed_wallet_balance": 201,
+                            "pending_change": 0,
+                            "pending_coin_removal_count": 3,
+                        },
+                    },
+                    post_block_balance_updates={
+                        "did_2": {
+                            "confirmed_wallet_balance": 201,
+                            "spendable_balance": 201,
+                            "max_send_amount": 201,
+                            "unspent_coin_count": 1,
+                            "pending_change": 0,
+                            "pending_coin_removal_count": -3,
+                        },
+                    },
+                ),
+                WalletStateTransition(
+                    pre_block_balance_updates={
+                        "did_2": {
+                            "unconfirmed_wallet_balance": 0,  # TODO: fix pre-block balances for recovery
+                            "spendable_balance": 0,
+                            "pending_change": 0,
+                            "max_send_amount": 0,
+                            "pending_coin_removal_count": 0,
+                            "set_remainder": True,
+                        },
+                    },
+                    post_block_balance_updates={
+                        "did_2": {
+                            "confirmed_wallet_balance": -201,
+                            "spendable_balance": -201,
+                            "max_send_amount": -201,
+                            "unspent_coin_count": -1,
+                            "set_remainder": True,
+                        },
+                    },
+                ),
+            ]
+        )
 
         for wallet in [did_wallet, did_wallet_2, did_wallet_3, did_wallet_4]:
             assert wallet.wallet_state_manager.wallets[wallet.id()] == wallet
 
-    @pytest.mark.parametrize(
-        "trusted",
-        [True, False],
-    )
+    @pytest.mark.limit_consensus_modes(allowed=[ConsensusMode.PLAIN], reason="irrelevant")
+    @pytest.mark.parametrize("wallet_environments", [{"num_environments": 1, "blocks_needed": [1]}], indirect=True)
     @pytest.mark.anyio
-    async def test_did_recovery_with_empty_set(self, self_hostname, two_wallet_nodes, trusted):
-        full_nodes, wallets, _ = two_wallet_nodes
-        full_node_api = full_nodes[0]
-        server_1 = full_node_api.server
-        wallet_node, server_2 = wallets[0]
-        wallet_node_2, server_3 = wallets[1]
-        wallet = wallet_node.wallet_state_manager.main_wallet
+    async def test_did_recovery_with_empty_set(self, wallet_environments: WalletTestFramework):
+        env_0 = wallet_environments.environments[0]
+        wallet_node_0 = env_0.node
+        wallet_0 = env_0.xch_wallet
 
-        ph = await wallet.get_new_puzzlehash()
-        if trusted:
-            wallet_node.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.config["trusted_peers"] = {}
-            wallet_node_2.config["trusted_peers"] = {}
-        await server_2.start_client(PeerInfo(self_hostname, server_1.get_port()), None)
-        await server_3.start_client(PeerInfo(self_hostname, server_1.get_port()), None)
+        env_0.wallet_aliases = {
+            "xch": 1,
+            "did": 2,
+        }
 
-        await full_node_api.farm_blocks_to_wallet(1, wallet)
+        ph = await wallet_0.get_new_puzzlehash()
 
-        async with wallet.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
+        async with wallet_0.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
             did_wallet: DIDWallet = await DIDWallet.create_new_did_wallet(
-                wallet_node.wallet_state_manager, wallet, uint64(101), action_scope
+                wallet_node_0.wallet_state_manager, wallet_0, uint64(101), action_scope
             )
 
-        await full_node_api.process_transaction_records(records=action_scope.side_effects.transactions)
-        await full_node_api.wait_for_wallets_synced(wallet_nodes=[wallet_node, wallet_node_2])
-
-        await time_out_assert(15, did_wallet.get_confirmed_balance, 101)
-        await time_out_assert(15, did_wallet.get_unconfirmed_balance, 101)
+        await wallet_environments.process_pending_states(
+            [
+                WalletStateTransition(
+                    pre_block_balance_updates={
+                        "xch": {
+                            "set_remainder": True,
+                        },
+                        "did": {
+                            "init": True,
+                            "unconfirmed_wallet_balance": 101,
+                            "pending_change": 101,
+                            "pending_coin_removal_count": 1,
+                        },
+                    },
+                    post_block_balance_updates={
+                        "xch": {
+                            "set_remainder": True,
+                        },
+                        "did": {
+                            "confirmed_wallet_balance": 101,
+                            "spendable_balance": 101,
+                            "max_send_amount": 101,
+                            "unspent_coin_count": 1,
+                            "pending_change": -101,
+                            "pending_coin_removal_count": -1,
+                        },
+                    },
+                ),
+            ]
+        )
         coin = await did_wallet.get_coin()
-        info = Program.to([])
+        info: list[tuple[bytes, bytes, int]] = []
         pubkey = (await did_wallet.wallet_state_manager.get_unused_derivation_record(did_wallet.wallet_info.id)).pubkey
         with pytest.raises(Exception):  # We expect a CLVM 80 error for this test
-            async with did_wallet.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=False) as action_scope:
+            async with did_wallet.wallet_state_manager.new_action_scope(
+                wallet_environments.tx_config, push=False
+            ) as action_scope:
                 await did_wallet.recovery_spend(
-                    coin, ph, info, pubkey, WalletSpendBundle([], AugSchemeMPL.aggregate([])), action_scope
+                    coin,
+                    ph,
+                    info,
+                    pubkey,
+                    WalletSpendBundle([], AugSchemeMPL.aggregate([])),
+                    action_scope,
                 )
 
-    @pytest.mark.parametrize(
-        "trusted",
-        [True, False],
-    )
+    @pytest.mark.limit_consensus_modes(allowed=[ConsensusMode.PLAIN], reason="irrelevant")
+    @pytest.mark.parametrize("wallet_environments", [{"num_environments": 1, "blocks_needed": [1]}], indirect=True)
     @pytest.mark.anyio
-    async def test_did_find_lost_did(self, self_hostname, two_wallet_nodes, trusted):
-        full_nodes, wallets, _ = two_wallet_nodes
-        full_node_api = full_nodes[0]
-        server_1 = full_node_api.server
-        wallet_node, server_2 = wallets[0]
-        wallet_node_2, server_3 = wallets[1]
-        wallet = wallet_node.wallet_state_manager.main_wallet
-        api_0 = WalletRpcApi(wallet_node)
-        if trusted:
-            wallet_node.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.config["trusted_peers"] = {}
-            wallet_node_2.config["trusted_peers"] = {}
-        await server_2.start_client(PeerInfo(self_hostname, server_1.get_port()), None)
-        await server_3.start_client(PeerInfo(self_hostname, server_1.get_port()), None)
-        await full_node_api.farm_blocks_to_wallet(1, wallet)
+    async def test_did_find_lost_did(self, wallet_environments: WalletTestFramework):
+        env_0 = wallet_environments.environments[0]
+        wallet_node_0 = env_0.node
+        wallet_0 = env_0.xch_wallet
+        api_0 = env_0.rpc_api
 
-        async with wallet.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
-            did_wallet: DIDWallet = await DIDWallet.create_new_did_wallet(
-                wallet_node.wallet_state_manager, wallet, uint64(101), action_scope
+        env_0.wallet_aliases = {
+            "xch": 1,
+            "did": 2,
+        }
+
+        async with wallet_0.wallet_state_manager.new_action_scope(
+            wallet_environments.tx_config, push=True
+        ) as action_scope:
+            did_wallet = await DIDWallet.create_new_did_wallet(
+                wallet_node_0.wallet_state_manager, wallet_0, uint64(101), action_scope
             )
-        await full_node_api.process_transaction_records(records=action_scope.side_effects.transactions)
-        await full_node_api.wait_for_wallets_synced(wallet_nodes=[wallet_node, wallet_node_2])
 
-        await time_out_assert(15, did_wallet.get_confirmed_balance, 101)
-        await time_out_assert(15, did_wallet.get_unconfirmed_balance, 101)
+        await wallet_environments.process_pending_states(
+            [
+                WalletStateTransition(
+                    pre_block_balance_updates={
+                        "xch": {
+                            "set_remainder": True,
+                        },
+                        "did": {
+                            "init": True,
+                            "unconfirmed_wallet_balance": 101,
+                            "pending_change": 101,
+                            "pending_coin_removal_count": 1,
+                        },
+                    },
+                    post_block_balance_updates={
+                        "xch": {
+                            "set_remainder": True,
+                        },
+                        "did": {
+                            "confirmed_wallet_balance": 101,
+                            "spendable_balance": 101,
+                            "max_send_amount": 101,
+                            "unspent_coin_count": 1,
+                            "pending_change": -101,
+                            "pending_coin_removal_count": -1,
+                        },
+                    },
+                ),
+            ]
+        )
+
         # Delete the coin and wallet
         coin = await did_wallet.get_coin()
-        await wallet_node.wallet_state_manager.coin_store.delete_coin_record(coin.name())
-        await time_out_assert(15, did_wallet.get_confirmed_balance, 0)
-        await wallet_node.wallet_state_manager.user_store.delete_wallet(did_wallet.wallet_info.id)
-        wallet_node.wallet_state_manager.wallets.pop(did_wallet.wallet_info.id)
-        assert len(wallet_node.wallet_state_manager.wallets) == 1
+        await wallet_node_0.wallet_state_manager.coin_store.delete_coin_record(coin.name())
+        await wallet_node_0.wallet_state_manager.delete_wallet(did_wallet.wallet_info.id)
+        wallet_node_0.wallet_state_manager.wallets.pop(did_wallet.wallet_info.id)
+        assert len(wallet_node_0.wallet_state_manager.wallets) == 1
         # Find lost DID
+        assert did_wallet.did_info.origin_coin is not None  # mypy
         resp = await api_0.did_find_lost_did({"coin_id": did_wallet.did_info.origin_coin.name().hex()})
         assert resp["success"]
         did_wallets = list(
             filter(
                 lambda w: (w.type == WalletType.DECENTRALIZED_ID),
-                await wallet_node.wallet_state_manager.get_all_wallet_info_entries(),
+                await wallet_node_0.wallet_state_manager.get_all_wallet_info_entries(),
             )
         )
-        did_wallet: Optional[DIDWallet] = wallet_node.wallet_state_manager.wallets[did_wallets[0].id]
-        await time_out_assert(15, did_wallet.get_confirmed_balance, 101)
-        await time_out_assert(15, did_wallet.get_unconfirmed_balance, 101)
+        did_wallet = wallet_node_0.wallet_state_manager.wallets[did_wallets[0].id]
+        env_0.wallet_aliases["did_found"] = did_wallets[0].id
+        await env_0.change_balances(
+            {
+                "did_found": {
+                    "init": True,
+                    "confirmed_wallet_balance": 101,
+                    "unconfirmed_wallet_balance": 101,
+                    "spendable_balance": 101,
+                    "max_send_amount": 101,
+                    "unspent_coin_count": 1,
+                }
+            }
+        )
+        await env_0.check_balances()
+
         # Spend DID
         recovery_list = [bytes32.fromhex(did_wallet.get_my_DID())]
         await did_wallet.update_recovery_list(recovery_list, uint64(1))
         assert did_wallet.did_info.backup_ids == recovery_list
-        async with did_wallet.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
+        async with did_wallet.wallet_state_manager.new_action_scope(
+            wallet_environments.tx_config, push=True
+        ) as action_scope:
             await did_wallet.create_update_spend(action_scope)
-        await full_node_api.process_transaction_records(records=action_scope.side_effects.transactions)
-        await full_node_api.wait_for_wallet_synced(wallet_node=wallet_node, timeout=20)
 
-        await time_out_assert(15, did_wallet.get_confirmed_balance, 101)
-        await time_out_assert(15, did_wallet.get_unconfirmed_balance, 101)
+        await wallet_environments.process_pending_states(
+            [
+                WalletStateTransition(
+                    pre_block_balance_updates={
+                        "did_found": {
+                            "spendable_balance": -101,
+                            "max_send_amount": -101,
+                            "pending_change": 101,
+                            "pending_coin_removal_count": 1,
+                        },
+                    },
+                    post_block_balance_updates={
+                        "did_found": {
+                            "spendable_balance": 101,
+                            "max_send_amount": 101,
+                            "pending_change": -101,
+                            "pending_coin_removal_count": -1,
+                        },
+                    },
+                ),
+            ]
+        )
+
         # Delete the coin and change inner puzzle
         coin = await did_wallet.get_coin()
-        await wallet_node.wallet_state_manager.coin_store.delete_coin_record(coin.name())
-        await time_out_assert(15, did_wallet.get_confirmed_balance, 0)
-        new_inner_puzzle = await did_wallet.get_new_did_innerpuz()
+        await wallet_node_0.wallet_state_manager.coin_store.delete_coin_record(coin.name())
+        new_inner_puzzle = await did_wallet.get_did_innerpuz(new=True)
         did_wallet.did_info = dataclasses.replace(did_wallet.did_info, current_inner=new_inner_puzzle)
         # Recovery the coin
+        assert did_wallet.did_info.origin_coin is not None  # mypy
         resp = await api_0.did_find_lost_did({"coin_id": did_wallet.did_info.origin_coin.name().hex()})
         assert resp["success"]
-        await time_out_assert(15, did_wallet.get_confirmed_balance, 101)
+        found_coin = await did_wallet.get_coin()
+        assert found_coin == coin
         assert did_wallet.did_info.current_inner != new_inner_puzzle
 
-    @pytest.mark.parametrize(
-        "trusted",
-        [True, False],
-    )
+    @pytest.mark.limit_consensus_modes(allowed=[ConsensusMode.PLAIN], reason="irrelevant")
+    @pytest.mark.parametrize("wallet_environments", [{"num_environments": 2, "blocks_needed": [1, 1]}], indirect=True)
     @pytest.mark.anyio
-    async def test_did_attest_after_recovery(self, self_hostname, two_wallet_nodes, trusted):
-        full_nodes, wallets, _ = two_wallet_nodes
-        full_node_api = full_nodes[0]
-        server_1 = full_node_api.server
-        wallet_node, server_2 = wallets[0]
-        wallet_node_2, server_3 = wallets[1]
-        wallet = wallet_node.wallet_state_manager.main_wallet
-        wallet2 = wallet_node_2.wallet_state_manager.main_wallet
-        if trusted:
-            wallet_node.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.config["trusted_peers"] = {}
-            wallet_node_2.config["trusted_peers"] = {}
-        await server_2.start_client(PeerInfo(self_hostname, server_1.get_port()), None)
-        await server_3.start_client(PeerInfo(self_hostname, server_1.get_port()), None)
-        await full_node_api.farm_blocks_to_wallet(1, wallet)
-        await full_node_api.farm_blocks_to_wallet(1, wallet2)
+    async def test_did_attest_after_recovery(self, wallet_environments: WalletTestFramework):
+        env_0 = wallet_environments.environments[0]
+        env_1 = wallet_environments.environments[1]
+        wallet_node_0 = env_0.node
+        wallet_node_1 = env_1.node
+        wallet_0 = env_0.xch_wallet
+        wallet_1 = env_1.xch_wallet
 
-        async with wallet.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
+        env_0.wallet_aliases = {
+            "xch": 1,
+            "did": 2,
+        }
+        env_1.wallet_aliases = {
+            "xch": 1,
+            "did": 2,
+        }
+
+        async with wallet_0.wallet_state_manager.new_action_scope(
+            wallet_environments.tx_config, push=True
+        ) as action_scope:
             did_wallet: DIDWallet = await DIDWallet.create_new_did_wallet(
-                wallet_node.wallet_state_manager, wallet, uint64(101), action_scope
+                wallet_node_0.wallet_state_manager, wallet_0, uint64(101), action_scope
             )
-        await full_node_api.process_transaction_records(records=action_scope.side_effects.transactions)
-        await full_node_api.wait_for_wallets_synced(wallet_nodes=[wallet_node, wallet_node_2])
+        await wallet_environments.process_pending_states(
+            [
+                WalletStateTransition(
+                    pre_block_balance_updates={
+                        "xch": {
+                            "set_remainder": True,
+                        },
+                        "did": {
+                            "init": True,
+                            "unconfirmed_wallet_balance": 101,
+                            "pending_change": 101,
+                            "pending_coin_removal_count": 1,
+                        },
+                    },
+                    post_block_balance_updates={
+                        "xch": {
+                            "set_remainder": True,
+                        },
+                        "did": {
+                            "confirmed_wallet_balance": 101,
+                            "spendable_balance": 101,
+                            "max_send_amount": 101,
+                            "unspent_coin_count": 1,
+                            "pending_change": -101,
+                            "pending_coin_removal_count": -1,
+                        },
+                    },
+                ),
+                WalletStateTransition(),
+            ]
+        )
 
         await time_out_assert(15, did_wallet.get_confirmed_balance, 101)
         await time_out_assert(15, did_wallet.get_unconfirmed_balance, 101)
-        recovery_list = [bytes.fromhex(did_wallet.get_my_DID())]
+        recovery_list = [bytes32.from_hexstr(did_wallet.get_my_DID())]
 
-        async with wallet2.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
+        async with wallet_1.wallet_state_manager.new_action_scope(
+            wallet_environments.tx_config, push=True
+        ) as action_scope:
             did_wallet_2: DIDWallet = await DIDWallet.create_new_did_wallet(
-                wallet_node_2.wallet_state_manager, wallet2, uint64(101), action_scope, recovery_list
+                wallet_node_1.wallet_state_manager, wallet_1, uint64(101), action_scope, recovery_list
             )
-        await full_node_api.process_transaction_records(records=action_scope.side_effects.transactions)
-        await full_node_api.wait_for_wallets_synced(wallet_nodes=[wallet_node, wallet_node_2])
-        await time_out_assert(25, did_wallet_2.get_confirmed_balance, 101)
-        await time_out_assert(25, did_wallet_2.get_unconfirmed_balance, 101)
+        await wallet_environments.process_pending_states(
+            [
+                WalletStateTransition(),
+                WalletStateTransition(
+                    pre_block_balance_updates={
+                        "xch": {
+                            "set_remainder": True,
+                        },
+                        "did": {
+                            "init": True,
+                            "unconfirmed_wallet_balance": 101,
+                            "pending_change": 101,
+                            "pending_coin_removal_count": 1,
+                        },
+                    },
+                    post_block_balance_updates={
+                        "xch": {
+                            "set_remainder": True,
+                        },
+                        "did": {
+                            "confirmed_wallet_balance": 101,
+                            "spendable_balance": 101,
+                            "max_send_amount": 101,
+                            "unspent_coin_count": 1,
+                            "pending_change": -101,
+                            "pending_coin_removal_count": -1,
+                        },
+                    },
+                ),
+            ]
+        )
+
         assert did_wallet_2.did_info.backup_ids == recovery_list
 
         # Update coin with new ID info
-        recovery_list = [bytes.fromhex(did_wallet_2.get_my_DID())]
+        recovery_list = [bytes32.from_hexstr(did_wallet_2.get_my_DID())]
         await did_wallet.update_recovery_list(recovery_list, uint64(1))
         assert did_wallet.did_info.backup_ids == recovery_list
-        async with did_wallet.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
+        async with did_wallet.wallet_state_manager.new_action_scope(
+            wallet_environments.tx_config, push=True
+        ) as action_scope:
             await did_wallet.create_update_spend(action_scope)
-        await full_node_api.process_transaction_records(records=action_scope.side_effects.transactions)
-        await full_node_api.wait_for_wallets_synced(wallet_nodes=[wallet_node, wallet_node_2])
 
-        await time_out_assert(15, did_wallet.get_confirmed_balance, 101)
-        await time_out_assert(15, did_wallet.get_unconfirmed_balance, 101)
+        await wallet_environments.process_pending_states(
+            [
+                WalletStateTransition(
+                    pre_block_balance_updates={
+                        "did": {
+                            "set_remainder": True,
+                        }
+                    },
+                    post_block_balance_updates={
+                        "did": {
+                            "set_remainder": True,
+                        }
+                    },
+                ),
+                WalletStateTransition(),
+            ]
+        )
 
         # DID Wallet 2 recovers into DID Wallet 3 with new innerpuz
         backup_data = did_wallet_2.create_backup()
 
-        async with wallet_node.wallet_state_manager.lock:
+        async with wallet_node_0.wallet_state_manager.lock:
             did_wallet_3 = await DIDWallet.create_new_did_wallet_from_recovery(
-                wallet_node.wallet_state_manager,
-                wallet,
+                wallet_node_0.wallet_state_manager,
+                wallet_0,
                 backup_data,
             )
-        new_ph = await did_wallet_3.get_new_did_inner_hash()
+        env_0.wallet_aliases["did_2"] = 3
+        new_ph = await did_wallet_3.get_did_inner_hash(new=True)
         coin = await did_wallet_2.get_coin()
         pubkey = (
             await did_wallet_3.wallet_state_manager.get_unused_derivation_record(did_wallet_3.wallet_info.id)
         ).pubkey
-        await time_out_assert(15, did_wallet.get_confirmed_balance, 101)
-        async with did_wallet.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
+
+        async with did_wallet.wallet_state_manager.new_action_scope(
+            wallet_environments.tx_config, push=True
+        ) as action_scope:
             message_spend_bundle, attest_data = await did_wallet.create_attestment(
                 coin.name(), new_ph, pubkey, action_scope
             )
-        await full_node_api.process_transaction_records(records=action_scope.side_effects.transactions)
-        await full_node_api.wait_for_wallets_synced(wallet_nodes=[wallet_node, wallet_node_2])
+        await wallet_environments.process_pending_states(
+            [
+                WalletStateTransition(
+                    pre_block_balance_updates={
+                        "did": {
+                            "set_remainder": True,
+                        },
+                        "did_2": {
+                            "init": True,
+                            "set_remainder": True,
+                        },
+                    },
+                    post_block_balance_updates={
+                        "did": {
+                            "set_remainder": True,
+                        },
+                        "did_2": {
+                            "set_remainder": True,
+                        },
+                    },
+                ),
+                WalletStateTransition(),
+            ]
+        )
 
         (
             info,
@@ -758,33 +1099,94 @@ class TestDIDWallet:
         ) = await did_wallet_3.load_attest_files_for_recovery_spend([attest_data])
         async with did_wallet_3.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
             await did_wallet_3.recovery_spend(coin, new_ph, info, pubkey, message_spend_bundle, action_scope)
-        await full_node_api.process_transaction_records(action_scope.side_effects.transactions)
-        await full_node_api.wait_for_wallets_synced(wallet_nodes=[wallet_node, wallet_node_2])
-
-        await time_out_assert(15, did_wallet_3.get_confirmed_balance, 101)
-        await time_out_assert(15, did_wallet_3.get_unconfirmed_balance, 101)
+        await wallet_environments.process_pending_states(
+            [
+                WalletStateTransition(
+                    pre_block_balance_updates={
+                        "did_2": {
+                            "unconfirmed_wallet_balance": 101,
+                            "set_remainder": True,
+                        },
+                    },
+                    post_block_balance_updates={
+                        "did_2": {
+                            "confirmed_wallet_balance": 101,
+                            "set_remainder": True,
+                        },
+                    },
+                ),
+                WalletStateTransition(
+                    pre_block_balance_updates={
+                        "did": {
+                            "unconfirmed_wallet_balance": 0,  # TODO: fix pre-block balances for recovery
+                            "set_remainder": True,
+                        },
+                    },
+                    post_block_balance_updates={
+                        "did": {
+                            "confirmed_wallet_balance": -101,
+                            "set_remainder": True,
+                        },
+                    },
+                ),
+            ]
+        )
 
         # DID Wallet 1 recovery spends into DID Wallet 4
         backup_data = did_wallet.create_backup()
 
-        async with wallet_node_2.wallet_state_manager.lock:
+        async with wallet_node_1.wallet_state_manager.lock:
             did_wallet_4 = await DIDWallet.create_new_did_wallet_from_recovery(
-                wallet_node_2.wallet_state_manager,
-                wallet2,
+                wallet_node_1.wallet_state_manager,
+                wallet_1,
                 backup_data,
             )
+        env_1.wallet_aliases["did_2"] = 3
         coin = await did_wallet.get_coin()
-        new_ph = await did_wallet_4.get_new_did_inner_hash()
+        new_ph = await did_wallet_4.get_did_inner_hash(new=True)
         pubkey = (
             await did_wallet_4.wallet_state_manager.get_unused_derivation_record(did_wallet_4.wallet_info.id)
         ).pubkey
-        async with did_wallet_3.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
+        async with did_wallet_3.wallet_state_manager.new_action_scope(
+            wallet_environments.tx_config, push=True
+        ) as action_scope:
             message_spend_bundle, attest1 = await did_wallet_3.create_attestment(
                 coin.name(), new_ph, pubkey, action_scope
             )
-        await full_node_api.process_transaction_records(records=action_scope.side_effects.transactions)
-        await full_node_api.wait_for_wallets_synced(wallet_nodes=[wallet_node, wallet_node_2])
-        await time_out_assert(15, wallet.get_pending_change_balance, 0)
+
+        await wallet_environments.process_pending_states(
+            [
+                WalletStateTransition(
+                    pre_block_balance_updates={
+                        "did_2": {
+                            "unconfirmed_wallet_balance": 0,
+                            "set_remainder": True,
+                        },
+                    },
+                    post_block_balance_updates={
+                        "did_2": {
+                            "confirmed_wallet_balance": 0,
+                            "set_remainder": True,
+                        },
+                    },
+                ),
+                WalletStateTransition(
+                    pre_block_balance_updates={
+                        "did_2": {
+                            "init": True,
+                            "set_remainder": True,
+                        },
+                    },
+                    post_block_balance_updates={
+                        "did_2": {
+                            "confirmed_wallet_balance": 0,
+                            "set_remainder": True,
+                        },
+                    },
+                ),
+            ]
+        )
+
         (
             test_info_list,
             test_message_spend_bundle,
@@ -794,27 +1196,186 @@ class TestDIDWallet:
                 coin, new_ph, test_info_list, pubkey, test_message_spend_bundle, action_scope
             )
 
-        await full_node_api.process_transaction_records(action_scope.side_effects.transactions)
-        await full_node_api.wait_for_wallets_synced(wallet_nodes=[wallet_node, wallet_node_2])
+        await wallet_environments.process_pending_states(
+            [
+                WalletStateTransition(
+                    pre_block_balance_updates={
+                        "did": {
+                            "set_remainder": True,
+                        },
+                    },
+                    post_block_balance_updates={
+                        "did": {
+                            "confirmed_wallet_balance": -101,
+                            "set_remainder": True,
+                        },
+                    },
+                ),
+                WalletStateTransition(
+                    pre_block_balance_updates={
+                        "did_2": {
+                            "set_remainder": True,
+                        },
+                    },
+                    post_block_balance_updates={
+                        "did_2": {
+                            "confirmed_wallet_balance": 101,
+                            "set_remainder": True,
+                        },
+                    },
+                ),
+            ]
+        )
+
+        for wallet in [did_wallet, did_wallet_3, did_wallet_4]:
+            assert wallet.wallet_state_manager.wallets[wallet.id()] == wallet
+
+    @pytest.mark.limit_consensus_modes(allowed=[ConsensusMode.PLAIN], reason="irrelevant")
+    @pytest.mark.parametrize("wallet_environments", [{"num_environments": 2, "blocks_needed": [1, 1]}], indirect=True)
+    @pytest.mark.parametrize(
+        "with_recovery",
+        [True, False],
+    )
+    @pytest.mark.anyio
+    async def test_did_transfer(self, wallet_environments: WalletTestFramework, with_recovery: bool):
+        env_0 = wallet_environments.environments[0]
+        env_1 = wallet_environments.environments[1]
+        wallet_node_0 = env_0.node
+        wallet_node_1 = env_1.node
+        wallet_0 = env_0.xch_wallet
+        wallet_1 = env_1.xch_wallet
+
+        env_0.wallet_aliases = {
+            "xch": 1,
+            "did": 2,
+        }
+        env_1.wallet_aliases = {
+            "xch": 1,
+            "did": 2,
+        }
+        ph = await wallet_0.get_new_puzzlehash()
+        fee = uint64(1000)
+
+        async with wallet_0.wallet_state_manager.new_action_scope(
+            wallet_environments.tx_config, push=True
+        ) as action_scope:
+            did_wallet_1: DIDWallet = await DIDWallet.create_new_did_wallet(
+                wallet_node_0.wallet_state_manager,
+                wallet_0,
+                uint64(101),
+                action_scope,
+                [ph],
+                uint64(1),
+                {"Twitter": "Test", "GitHub": "测试"},
+                fee=fee,
+            )
+        assert did_wallet_1.get_name() == "Profile 1"
+
+        await wallet_environments.process_pending_states(
+            [
+                WalletStateTransition(
+                    pre_block_balance_updates={
+                        "xch": {
+                            "set_remainder": True,
+                        },
+                        "did": {
+                            "init": True,
+                            "unconfirmed_wallet_balance": 101,
+                            "pending_change": 101,
+                            "pending_coin_removal_count": 1,
+                        },
+                    },
+                    post_block_balance_updates={
+                        "xch": {
+                            "set_remainder": True,
+                        },
+                        "did": {
+                            "confirmed_wallet_balance": 101,
+                            "spendable_balance": 101,
+                            "max_send_amount": 101,
+                            "unspent_coin_count": 1,
+                            "pending_change": -101,
+                            "pending_coin_removal_count": -1,
+                        },
+                    },
+                ),
+                WalletStateTransition(),
+            ]
+        )
+
+        # Transfer DID
+        new_puzhash = await wallet_1.get_new_puzzlehash()
+        async with did_wallet_1.wallet_state_manager.new_action_scope(
+            wallet_environments.tx_config, push=True
+        ) as action_scope:
+            await did_wallet_1.transfer_did(new_puzhash, fee, with_recovery, action_scope)
+
+        await wallet_environments.process_pending_states(
+            [
+                WalletStateTransition(
+                    pre_block_balance_updates={
+                        "xch": {
+                            "set_remainder": True,
+                        },
+                        "did": {
+                            "unconfirmed_wallet_balance": -101,
+                            "set_remainder": True,
+                        },
+                    },
+                    post_block_balance_updates={
+                        "xch": {
+                            "set_remainder": True,
+                        },
+                    },
+                ),
+                WalletStateTransition(
+                    pre_block_balance_updates={},
+                    post_block_balance_updates={
+                        "did": {
+                            "init": True,
+                            "confirmed_wallet_balance": 101,
+                            "set_remainder": True,
+                        },
+                    },
+                ),
+            ]
+        )
 
-        await time_out_assert(15, did_wallet_4.get_confirmed_balance, 101)
-        await time_out_assert(15, did_wallet_4.get_unconfirmed_balance, 101)
-        await time_out_assert(15, did_wallet.get_confirmed_balance, 0)
-        await time_out_assert(15, did_wallet.get_unconfirmed_balance, 0)
+        # Get the new DID wallet
+        did_wallets = list(
+            filter(
+                lambda w: (w.type == WalletType.DECENTRALIZED_ID),
+                await wallet_node_1.wallet_state_manager.get_all_wallet_info_entries(),
+            )
+        )
+        did_wallet_2 = wallet_node_1.wallet_state_manager.wallets[did_wallets[0].id]
+        assert isinstance(did_wallet_2, DIDWallet)  # mypy
+        assert len(wallet_node_0.wallet_state_manager.wallets) == 1
+        assert did_wallet_1.did_info.origin_coin == did_wallet_2.did_info.origin_coin
+        if with_recovery:
+            assert did_wallet_1.did_info.backup_ids[0] == did_wallet_2.did_info.backup_ids[0]
+            assert did_wallet_1.did_info.num_of_backup_ids_needed == did_wallet_2.did_info.num_of_backup_ids_needed
+        metadata = json.loads(did_wallet_2.did_info.metadata)
+        assert metadata["Twitter"] == "Test"
+        assert metadata["GitHub"] == "测试"
 
-        for wallet in [did_wallet, did_wallet_3, did_wallet_4]:
-            assert wallet.wallet_state_manager.wallets[wallet.id()] == wallet
+        # Test match_hinted_coin
+        assert await did_wallet_2.match_hinted_coin(
+            await did_wallet_2.get_coin(),
+            new_puzhash,
+        )
 
-    @pytest.mark.parametrize(
-        "with_recovery",
-        [True, False],
-    )
     @pytest.mark.parametrize(
         "trusted",
         [True, False],
     )
     @pytest.mark.anyio
-    async def test_did_transfer(self, self_hostname, two_wallet_nodes, with_recovery, trusted):
+    async def test_did_auto_transfer_limit(
+        self,
+        self_hostname: str,
+        two_wallet_nodes: tuple[list[FullNodeSimulator], list[tuple[WalletNode, ChiaServer]], BlockTools],
+        trusted: bool,
+    ) -> None:
         fee = uint64(1000)
         full_nodes, wallets, _ = two_wallet_nodes
         full_node_api = full_nodes[0]
@@ -823,6 +1384,7 @@ class TestDIDWallet:
         wallet_node_2, server_3 = wallets[1]
         wallet = wallet_node.wallet_state_manager.main_wallet
         wallet2 = wallet_node_2.wallet_state_manager.main_wallet
+        api_1 = WalletRpcApi(wallet_node_2)
         ph = await wallet.get_new_puzzlehash()
 
         if trusted:
@@ -840,138 +1402,235 @@ class TestDIDWallet:
         await server_3.start_client(PeerInfo(self_hostname, server_1.get_port()), None)
         await full_node_api.farm_blocks_to_wallet(1, wallet)
 
-        async with wallet.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
-            did_wallet_1: DIDWallet = await DIDWallet.create_new_did_wallet(
-                wallet_node.wallet_state_manager,
-                wallet,
+        # Check that we cap out at 10 DID Wallets automatically created upon transfer received
+        for i in range(0, 14):
+            async with wallet.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
+                did_wallet_1: DIDWallet = await DIDWallet.create_new_did_wallet(
+                    wallet_node.wallet_state_manager,
+                    wallet,
+                    uint64(101),
+                    action_scope,
+                    [bytes32(bytes(ph))],
+                    uint64(1),
+                    {"Twitter": "Test", "GitHub": "测试"},
+                    fee=fee,
+                )
+            assert did_wallet_1.get_name() == "Profile 1"
+            await full_node_api.process_transaction_records(records=action_scope.side_effects.transactions)
+            await full_node_api.wait_for_wallets_synced(wallet_nodes=[wallet_node, wallet_node_2])
+            await time_out_assert(15, did_wallet_1.get_confirmed_balance, 101)
+            await time_out_assert(15, did_wallet_1.get_unconfirmed_balance, 101)
+            # Transfer DID
+            assert did_wallet_1.did_info.origin_coin is not None
+            origin_coin = did_wallet_1.did_info.origin_coin
+            new_puzhash = await wallet2.get_new_puzzlehash()
+            async with did_wallet_1.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
+                await did_wallet_1.transfer_did(new_puzhash, fee, False, action_scope)
+            await full_node_api.process_transaction_records(records=action_scope.side_effects.transactions)
+            await full_node_api.wait_for_wallets_synced(wallet_nodes=[wallet_node, wallet_node_2])
+            # Check if the DID wallet is created in the wallet2
+
+            await time_out_assert(
+                30, get_wallet_num, min(2 + i, 11), wallet_node_2.wallet_state_manager
+            )  # check we haven't made more than 10 DID wallets
+            await time_out_assert(30, get_wallet_num, 1, wallet_node.wallet_state_manager)
+        # Get the new DID wallets
+        did_wallets = list(
+            filter(
+                lambda w: (w.type == WalletType.DECENTRALIZED_ID),
+                await wallet_node_2.wallet_state_manager.get_all_wallet_info_entries(),
+            )
+        )
+
+        assert len(did_wallets) == 10
+        # Test we can use the DID
+        did_wallet_10 = wallet_node_2.wallet_state_manager.get_wallet(
+            id=uint32(did_wallets[9].id), required_type=DIDWallet
+        )
+        # Delete the coin and change inner puzzle
+        coin = await did_wallet_10.get_coin()
+        # origin_coin = did_wallet_10.did_info.origin_coin
+        backup_data = did_wallet_10.create_backup()
+        await wallet_node_2.wallet_state_manager.coin_store.delete_coin_record(coin.name())
+        await time_out_assert(15, did_wallet_10.get_confirmed_balance, 0)
+        await wallet_node_2.wallet_state_manager.user_store.delete_wallet(did_wallet_10.wallet_info.id)
+        wallet_node_2.wallet_state_manager.wallets.pop(did_wallet_10.wallet_info.id)
+        # Recover the coin
+        async with wallet_node_2.wallet_state_manager.lock:
+            did_wallet_10 = await DIDWallet.create_new_did_wallet_from_recovery(
+                wallet_node_2.wallet_state_manager,
+                wallet2,
+                backup_data,
+            )
+        assert did_wallet_10.did_info.origin_coin is not None
+        resp = await api_1.did_find_lost_did({"coin_id": did_wallet_10.did_info.origin_coin.name().hex()})
+        assert resp["success"]
+        await time_out_assert(15, did_wallet_10.get_confirmed_balance, 101)
+        await time_out_assert(15, did_wallet_10.get_unconfirmed_balance, 101)
+
+        # Check we can recover an auto-discarded DID
+        did_wallet_9 = wallet_node_2.wallet_state_manager.get_wallet(
+            id=uint32(did_wallets[8].id), required_type=DIDWallet
+        )
+        # Delete the coin and wallet to make space for a auto-discarded DID
+        coin = await did_wallet_9.get_coin()
+        await wallet_node_2.wallet_state_manager.coin_store.delete_coin_record(coin.name())
+        await time_out_assert(15, did_wallet_9.get_confirmed_balance, 0)
+        await wallet_node_2.wallet_state_manager.user_store.delete_wallet(did_wallet_9.wallet_info.id)
+        wallet_node_2.wallet_state_manager.wallets.pop(did_wallet_9.wallet_info.id)
+
+        did_wallets = list(
+            filter(
+                lambda w: (w.type == WalletType.DECENTRALIZED_ID),
+                await wallet_node_2.wallet_state_manager.get_all_wallet_info_entries(),
+            )
+        )
+        assert len(did_wallets) == 9
+
+        # Try and find lost coin
+        resp = await api_1.did_find_lost_did({"coin_id": origin_coin.name().hex()})
+        did_wallets = list(
+            filter(
+                lambda w: (w.type == WalletType.DECENTRALIZED_ID),
+                await wallet_node_2.wallet_state_manager.get_all_wallet_info_entries(),
+            )
+        )
+        assert len(did_wallets) == 10
+
+        # Check we can still manually add new DIDs while at cap
+        await full_node_api.farm_blocks_to_wallet(1, wallet2)
+        ph = await wallet2.get_new_puzzlehash()
+        async with wallet2.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
+            did_wallet_11: DIDWallet = await DIDWallet.create_new_did_wallet(
+                wallet_node_2.wallet_state_manager,
+                wallet2,
                 uint64(101),
                 action_scope,
-                [bytes(ph)],
+                [bytes32(bytes(ph))],
                 uint64(1),
                 {"Twitter": "Test", "GitHub": "测试"},
                 fee=fee,
             )
-        assert did_wallet_1.get_name() == "Profile 1"
         await full_node_api.process_transaction_records(records=action_scope.side_effects.transactions)
         await full_node_api.wait_for_wallets_synced(wallet_nodes=[wallet_node, wallet_node_2])
-        await time_out_assert(15, did_wallet_1.get_confirmed_balance, 101)
-        await time_out_assert(15, did_wallet_1.get_unconfirmed_balance, 101)
-        # Transfer DID
-        new_puzhash = await wallet2.get_new_puzzlehash()
-        async with did_wallet_1.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
-            await did_wallet_1.transfer_did(new_puzhash, fee, with_recovery, action_scope)
-        await full_node_api.process_transaction_records(records=action_scope.side_effects.transactions)
-        await full_node_api.wait_for_wallets_synced(wallet_nodes=[wallet_node, wallet_node_2])
-        # Check if the DID wallet is created in the wallet2
+        await time_out_assert(15, did_wallet_11.get_confirmed_balance, 101)
+        await time_out_assert(15, did_wallet_11.get_unconfirmed_balance, 101)
 
-        await time_out_assert(30, get_wallet_num, 2, wallet_node_2.wallet_state_manager)
-        await time_out_assert(30, get_wallet_num, 1, wallet_node.wallet_state_manager)
-        # Get the new DID wallet
         did_wallets = list(
             filter(
                 lambda w: (w.type == WalletType.DECENTRALIZED_ID),
                 await wallet_node_2.wallet_state_manager.get_all_wallet_info_entries(),
             )
         )
-        did_wallet_2: Optional[DIDWallet] = wallet_node_2.wallet_state_manager.wallets[did_wallets[0].id]
-        assert len(wallet_node.wallet_state_manager.wallets) == 1
-        assert did_wallet_1.did_info.origin_coin == did_wallet_2.did_info.origin_coin
-        if with_recovery:
-            assert did_wallet_1.did_info.backup_ids[0] == did_wallet_2.did_info.backup_ids[0]
-            assert did_wallet_1.did_info.num_of_backup_ids_needed == did_wallet_2.did_info.num_of_backup_ids_needed
-        metadata = json.loads(did_wallet_2.did_info.metadata)
-        assert metadata["Twitter"] == "Test"
-        assert metadata["GitHub"] == "测试"
-
-        # Test match_hinted_coin
-        assert await did_wallet_2.match_hinted_coin(
-            await did_wallet_2.get_coin(),
-            new_puzhash,
-        )
+        assert len(did_wallets) == 11
 
-    @pytest.mark.parametrize(
-        "trusted",
-        [True, False],
-    )
+    @pytest.mark.limit_consensus_modes(allowed=[ConsensusMode.PLAIN], reason="irrelevant")
+    @pytest.mark.parametrize("wallet_environments", [{"num_environments": 1, "blocks_needed": [1]}], indirect=True)
     @pytest.mark.anyio
-    async def test_update_recovery_list(self, self_hostname, two_wallet_nodes, trusted):
-        full_nodes, wallets, _ = two_wallet_nodes
-        full_node_api = full_nodes[0]
-        server_1 = full_node_api.server
-        wallet_node, server_2 = wallets[0]
-        wallet_node_2, server_3 = wallets[1]
-        wallet = wallet_node.wallet_state_manager.main_wallet
-        ph = await wallet.get_new_puzzlehash()
+    async def test_update_recovery_list(self, wallet_environments: WalletTestFramework):
+        env = wallet_environments.environments[0]
+        wallet_node = env.node
+        wallet = env.xch_wallet
 
-        if trusted:
-            wallet_node.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.config["trusted_peers"] = {}
-            wallet_node_2.config["trusted_peers"] = {}
+        env.wallet_aliases = {
+            "xch": 1,
+            "did": 2,
+        }
 
-        await server_2.start_client(PeerInfo(self_hostname, server_1.get_port()), None)
-        await server_3.start_client(PeerInfo(self_hostname, server_1.get_port()), None)
-        await full_node_api.farm_blocks_to_wallet(1, wallet)
+        ph = await wallet.get_new_puzzlehash()
 
-        async with wallet.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
+        async with wallet.wallet_state_manager.new_action_scope(
+            wallet_environments.tx_config, push=True
+        ) as action_scope:
             did_wallet_1: DIDWallet = await DIDWallet.create_new_did_wallet(
                 wallet_node.wallet_state_manager, wallet, uint64(101), action_scope, []
             )
-        await full_node_api.process_transaction_records(records=action_scope.side_effects.transactions)
-        await full_node_api.wait_for_wallets_synced(wallet_nodes=[wallet_node, wallet_node_2])
-        await time_out_assert(15, did_wallet_1.get_confirmed_balance, 101)
-        await time_out_assert(15, did_wallet_1.get_unconfirmed_balance, 101)
-        await did_wallet_1.update_recovery_list([bytes(ph)], 1)
-        async with did_wallet_1.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
+
+        await wallet_environments.process_pending_states(
+            [
+                WalletStateTransition(
+                    pre_block_balance_updates={
+                        "xch": {
+                            "set_remainder": True,
+                        },
+                        "did": {
+                            "init": True,
+                            "unconfirmed_wallet_balance": 101,
+                            "pending_change": 101,
+                            "pending_coin_removal_count": 1,
+                        },
+                    },
+                    post_block_balance_updates={
+                        "xch": {
+                            "set_remainder": True,
+                        },
+                        "did": {
+                            "confirmed_wallet_balance": 101,
+                            "spendable_balance": 101,
+                            "max_send_amount": 101,
+                            "unspent_coin_count": 1,
+                            "pending_change": -101,
+                            "pending_coin_removal_count": -1,
+                        },
+                    },
+                ),
+            ]
+        )
+        await did_wallet_1.update_recovery_list([ph], uint64(1))
+        async with did_wallet_1.wallet_state_manager.new_action_scope(
+            wallet_environments.tx_config, push=True
+        ) as action_scope:
             await did_wallet_1.create_update_spend(action_scope)
-        await full_node_api.process_transaction_records(records=action_scope.side_effects.transactions)
-        await full_node_api.wait_for_wallets_synced(wallet_nodes=[wallet_node, wallet_node_2])
-        await time_out_assert(15, did_wallet_1.get_confirmed_balance, 101)
-        await time_out_assert(15, did_wallet_1.get_unconfirmed_balance, 101)
+
+        await wallet_environments.process_pending_states(
+            [
+                WalletStateTransition(
+                    pre_block_balance_updates={
+                        "did": {
+                            "set_remainder": True,
+                        },
+                    },
+                    post_block_balance_updates={
+                        "did": {
+                            "set_remainder": True,
+                        },
+                    },
+                ),
+            ]
+        )
         assert did_wallet_1.did_info.backup_ids[0] == bytes(ph)
         assert did_wallet_1.did_info.num_of_backup_ids_needed == 1
 
-    @pytest.mark.parametrize(
-        "trusted",
-        [True, False],
-    )
+    @pytest.mark.limit_consensus_modes(allowed=[ConsensusMode.PLAIN], reason="irrelevant")
+    @pytest.mark.parametrize("wallet_environments", [{"num_environments": 2, "blocks_needed": [1, 1]}], indirect=True)
     @pytest.mark.anyio
-    async def test_get_info(self, self_hostname, two_wallet_nodes, trusted):
-        fee = uint64(1000)
-        full_nodes, wallets, _ = two_wallet_nodes
-        full_node_api = full_nodes[0]
-        server_1 = full_node_api.server
-        wallet_node, server_2 = wallets[0]
-        wallet_node_2, server_3 = wallets[1]
-        wallet = wallet_node.wallet_state_manager.main_wallet
-        wallet1 = wallet_node_2.wallet_state_manager.main_wallet
-        ph1 = await wallet1.get_new_puzzlehash()
-        api_0 = WalletRpcApi(wallet_node)
-        if trusted:
-            wallet_node.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.config["trusted_peers"] = {}
-            wallet_node_2.config["trusted_peers"] = {}
+    async def test_get_info(self, wallet_environments: WalletTestFramework):
+        env_0 = wallet_environments.environments[0]
+        env_1 = wallet_environments.environments[1]
+        wallet_node_0 = env_0.node
+        wallet_0 = env_0.xch_wallet
+        wallet_1 = env_1.xch_wallet
+        api_0 = env_0.rpc_api
 
-        await server_2.start_client(PeerInfo(self_hostname, server_1.get_port()), None)
-        await server_3.start_client(PeerInfo(self_hostname, server_1.get_port()), None)
-        await full_node_api.farm_blocks_to_wallet(count=2, wallet=wallet)
+        env_0.wallet_aliases = {
+            "xch": 1,
+            "did": 2,
+        }
+        env_1.wallet_aliases = {
+            "xch": 1,
+            "did": 2,
+        }
+
+        fee = uint64(1000)
         did_amount = uint64(101)
+        ph_1 = await wallet_1.get_new_puzzlehash()
 
-        async with wallet.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
+        async with wallet_0.wallet_state_manager.new_action_scope(
+            wallet_environments.tx_config, push=True
+        ) as action_scope:
             did_wallet_1: DIDWallet = await DIDWallet.create_new_did_wallet(
-                wallet_node.wallet_state_manager,
-                wallet,
+                wallet_node_0.wallet_state_manager,
+                wallet_0,
                 did_amount,
                 action_scope,
                 [],
@@ -979,14 +1638,42 @@ class TestDIDWallet:
                 fee=fee,
             )
 
-        await full_node_api.process_transaction_records(records=action_scope.side_effects.transactions)
-        await full_node_api.wait_for_wallet_synced(wallet_node=wallet_node, timeout=15)
-
-        assert await did_wallet_1.get_confirmed_balance() == did_amount
-        assert await did_wallet_1.get_unconfirmed_balance() == did_amount
+        await wallet_environments.process_pending_states(
+            [
+                WalletStateTransition(
+                    pre_block_balance_updates={
+                        "xch": {
+                            "set_remainder": True,
+                        },
+                        "did": {
+                            "init": True,
+                            "unconfirmed_wallet_balance": did_amount,
+                            "pending_change": did_amount,
+                            "pending_coin_removal_count": 1,
+                        },
+                    },
+                    post_block_balance_updates={
+                        "xch": {
+                            "set_remainder": True,
+                        },
+                        "did": {
+                            "confirmed_wallet_balance": did_amount,
+                            "spendable_balance": did_amount,
+                            "max_send_amount": did_amount,
+                            "unspent_coin_count": 1,
+                            "pending_change": -did_amount,
+                            "pending_coin_removal_count": -1,
+                        },
+                    },
+                ),
+                WalletStateTransition(),
+            ]
+        )
+        assert did_wallet_1.did_info.origin_coin is not None  # mypy
         response = await api_0.did_get_info({"coin_id": did_wallet_1.did_info.origin_coin.name().hex()})
         assert response["did_id"] == encode_puzzle_hash(did_wallet_1.did_info.origin_coin.name(), AddressType.DID.value)
         assert response["launcher_id"] == did_wallet_1.did_info.origin_coin.name().hex()
+        assert did_wallet_1.did_info.current_inner is not None  # mypy
         assert response["full_puzzle"].to_program() == create_singleton_puzzle(
             did_wallet_1.did_info.current_inner, did_wallet_1.did_info.origin_coin.name()
         )
@@ -997,8 +1684,10 @@ class TestDIDWallet:
         assert decode_puzzle_hash(response["p2_address"]).hex() == response["hints"][0]
 
         # Test non-singleton coin
-        async with wallet.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
-            coin = (await wallet.select_coins(uint64(1), action_scope)).pop()
+        async with wallet_0.wallet_state_manager.new_action_scope(
+            wallet_environments.tx_config, push=True
+        ) as action_scope:
+            coin = (await wallet_0.select_coins(uint64(1), action_scope)).pop()
         assert coin.amount % 2 == 1
         coin_id = coin.name()
         response = await api_0.did_get_info({"coin_id": coin_id.hex()})
@@ -1006,60 +1695,112 @@ class TestDIDWallet:
 
         # Test multiple odd coins
         odd_amount = uint64(1)
-        async with wallet.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
+        async with wallet_0.wallet_state_manager.new_action_scope(
+            wallet_environments.tx_config, push=True
+        ) as action_scope:
             async with action_scope.use() as interface:
                 interface.side_effects.selected_coins.append(coin)
-            coin_1 = (await wallet.select_coins(odd_amount, action_scope)).pop()
+            coin_1 = (await wallet_0.select_coins(odd_amount, action_scope)).pop()
         assert coin_1.amount % 2 == 0
-        async with wallet.wallet_state_manager.new_action_scope(
-            DEFAULT_TX_CONFIG.override(excluded_coin_ids=[coin_id]), push=True
+        async with wallet_0.wallet_state_manager.new_action_scope(
+            wallet_environments.tx_config.override(excluded_coin_ids=[coin_id]), push=True
         ) as action_scope:
-            await wallet.generate_signed_transaction(odd_amount, ph1, action_scope, fee)
-        await full_node_api.process_transaction_records(records=action_scope.side_effects.transactions)
-        await full_node_api.wait_for_wallet_synced(wallet_node=wallet_node_2, timeout=15)
+            await wallet_0.generate_signed_transaction(odd_amount, ph_1, action_scope, fee)
+
+        await wallet_environments.process_pending_states(
+            [
+                WalletStateTransition(
+                    pre_block_balance_updates={
+                        "xch": {
+                            "unconfirmed_wallet_balance": -odd_amount - fee,
+                            "set_remainder": True,
+                        },
+                        "did": {
+                            "set_remainder": True,
+                        },
+                    },
+                    post_block_balance_updates={
+                        "xch": {
+                            "confirmed_wallet_balance": -odd_amount - fee,
+                            "set_remainder": True,
+                        },
+                        "did": {
+                            "set_remainder": True,
+                        },
+                    },
+                ),
+                WalletStateTransition(
+                    pre_block_balance_updates={
+                        "xch": {
+                            "unconfirmed_wallet_balance": 0,
+                            "set_remainder": True,
+                        }
+                    },
+                    post_block_balance_updates={
+                        "xch": {
+                            "confirmed_wallet_balance": odd_amount,
+                            "set_remainder": True,
+                        }
+                    },
+                ),
+            ]
+        )
 
-        assert await wallet1.get_confirmed_balance() == odd_amount
         with pytest.raises(ValueError):
             await api_0.did_get_info({"coin_id": coin_1.name().hex()})
 
-    @pytest.mark.parametrize(
-        "trusted",
-        [True, False],
-    )
+    @pytest.mark.limit_consensus_modes(allowed=[ConsensusMode.PLAIN], reason="irrelevant")
+    @pytest.mark.parametrize("wallet_environments", [{"num_environments": 1, "blocks_needed": [1]}], indirect=True)
     @pytest.mark.anyio
-    async def test_message_spend(self, self_hostname, two_wallet_nodes, trusted):
-        fee = uint64(1000)
-        full_nodes, wallets, _ = two_wallet_nodes
-        full_node_api = full_nodes[0]
-        server_1 = full_node_api.server
-        wallet_node, server_2 = wallets[0]
-        wallet_node_2, server_3 = wallets[1]
-        wallet = wallet_node.wallet_state_manager.main_wallet
-        api_0 = WalletRpcApi(wallet_node)
-        if trusted:
-            wallet_node.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.config["trusted_peers"] = {}
-            wallet_node_2.config["trusted_peers"] = {}
+    async def test_message_spend(self, wallet_environments: WalletTestFramework):
+        env = wallet_environments.environments[0]
+        wallet_node = env.node
+        wallet = env.xch_wallet
+        api_0 = env.rpc_api
 
-        await server_2.start_client(PeerInfo(self_hostname, server_1.get_port()), None)
-        await server_3.start_client(PeerInfo(self_hostname, server_1.get_port()), None)
+        env.wallet_aliases = {
+            "xch": 1,
+            "did": 2,
+        }
 
-        await full_node_api.farm_blocks_to_wallet(1, wallet)
+        fee = uint64(1000)
 
-        async with wallet.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
+        async with wallet.wallet_state_manager.new_action_scope(
+            wallet_environments.tx_config, push=True
+        ) as action_scope:
             did_wallet_1: DIDWallet = await DIDWallet.create_new_did_wallet(
                 wallet_node.wallet_state_manager, wallet, uint64(101), action_scope, [], fee=fee
             )
-        await full_node_api.process_transaction_records(records=action_scope.side_effects.transactions)
-        await full_node_api.wait_for_wallets_synced(wallet_nodes=[wallet_node, wallet_node_2])
-        await time_out_assert(15, did_wallet_1.get_confirmed_balance, 101)
-        await time_out_assert(15, did_wallet_1.get_unconfirmed_balance, 101)
+        await wallet_environments.process_pending_states(
+            [
+                WalletStateTransition(
+                    pre_block_balance_updates={
+                        "xch": {
+                            "set_remainder": True,
+                        },
+                        "did": {
+                            "init": True,
+                            "unconfirmed_wallet_balance": 101,
+                            "pending_change": 101,
+                            "pending_coin_removal_count": 1,
+                        },
+                    },
+                    post_block_balance_updates={
+                        "xch": {
+                            "set_remainder": True,
+                        },
+                        "did": {
+                            "confirmed_wallet_balance": 101,
+                            "spendable_balance": 101,
+                            "max_send_amount": 101,
+                            "unspent_coin_count": 1,
+                            "pending_change": -101,
+                            "pending_coin_removal_count": -1,
+                        },
+                    },
+                ),
+            ]
+        )
         response = await api_0.did_message_spend(
             {"wallet_id": did_wallet_1.wallet_id, "coin_announcements": ["0abc"], "puzzle_announcements": ["0def"]}
         )
@@ -1073,125 +1814,174 @@ class TestDIDWallet:
         assert len(conditions[ConditionOpcode.CREATE_PUZZLE_ANNOUNCEMENT]) == 1
         assert conditions[ConditionOpcode.CREATE_PUZZLE_ANNOUNCEMENT][0].vars[0].hex() == "0def"
 
-    @pytest.mark.parametrize(
-        "trusted",
-        [True, False],
-    )
+    @pytest.mark.limit_consensus_modes(allowed=[ConsensusMode.PLAIN], reason="irrelevant")
+    @pytest.mark.parametrize("wallet_environments", [{"num_environments": 1, "blocks_needed": [1]}], indirect=True)
     @pytest.mark.anyio
-    async def test_update_metadata(self, self_hostname, two_wallet_nodes, trusted):
-        fee = uint64(1000)
-        full_nodes, wallets, _ = two_wallet_nodes
-        full_node_api = full_nodes[0]
-        server_1 = full_node_api.server
-        wallet_node, server_2 = wallets[0]
-        wallet_node_2, server_3 = wallets[1]
-        wallet = wallet_node.wallet_state_manager.main_wallet
-        if trusted:
-            wallet_node.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.config["trusted_peers"] = {}
-            wallet_node_2.config["trusted_peers"] = {}
+    async def test_update_metadata(self, wallet_environments: WalletTestFramework):
+        env = wallet_environments.environments[0]
+        wallet_node = env.node
+        wallet = env.xch_wallet
 
-        await server_2.start_client(PeerInfo(self_hostname, server_1.get_port()), None)
-        await server_3.start_client(PeerInfo(self_hostname, server_1.get_port()), None)
-        expected_confirmed_balance = await full_node_api.farm_blocks_to_wallet(count=2, wallet=wallet)
+        env.wallet_aliases = {
+            "xch": 1,
+            "did": 2,
+        }
+
+        fee = uint64(1000)
         did_amount = uint64(101)
 
-        async with wallet.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
+        async with wallet.wallet_state_manager.new_action_scope(
+            wallet_environments.tx_config, push=True
+        ) as action_scope:
             did_wallet_1: DIDWallet = await DIDWallet.create_new_did_wallet(
                 wallet_node.wallet_state_manager, wallet, did_amount, action_scope, [], fee=fee
             )
-        await full_node_api.process_transaction_records(records=action_scope.side_effects.transactions)
-        await full_node_api.wait_for_wallet_synced(wallet_node=wallet_node, timeout=15)
 
-        expected_confirmed_balance -= did_amount + fee
-        assert await did_wallet_1.get_confirmed_balance() == did_amount
-        assert await did_wallet_1.get_unconfirmed_balance() == did_amount
-        assert await wallet.get_confirmed_balance() == expected_confirmed_balance
-        assert await wallet.get_unconfirmed_balance() == expected_confirmed_balance
+        await wallet_environments.process_pending_states(
+            [
+                WalletStateTransition(
+                    pre_block_balance_updates={
+                        "xch": {
+                            "set_remainder": True,
+                        },
+                        "did": {
+                            "init": True,
+                            "unconfirmed_wallet_balance": 101,
+                            "pending_change": 101,
+                            "pending_coin_removal_count": 1,
+                        },
+                    },
+                    post_block_balance_updates={
+                        "xch": {
+                            "set_remainder": True,
+                        },
+                        "did": {
+                            "confirmed_wallet_balance": 101,
+                            "spendable_balance": 101,
+                            "max_send_amount": 101,
+                            "unspent_coin_count": 1,
+                            "pending_change": -101,
+                            "pending_coin_removal_count": -1,
+                        },
+                    },
+                ),
+            ]
+        )
+
+        assert did_wallet_1.did_info.current_inner is not None  # mypy
         puzhash = did_wallet_1.did_info.current_inner.get_tree_hash()
         parent_num = get_parent_num(did_wallet_1)
 
         bad_metadata = {"Twitter": {"url": "http://www.twitter.com"}}
         with pytest.raises(ValueError) as e:
-            await did_wallet_1.update_metadata(bad_metadata)
+            await did_wallet_1.update_metadata(bad_metadata)  # type: ignore
         assert e.match("Metadata key value pairs must be strings.")
 
         metadata = {}
         metadata["Twitter"] = "http://www.twitter.com"
         await did_wallet_1.update_metadata(metadata)
-        async with did_wallet_1.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
+        async with did_wallet_1.wallet_state_manager.new_action_scope(
+            wallet_environments.tx_config, push=True
+        ) as action_scope:
             await did_wallet_1.create_update_spend(action_scope, fee)
-        await full_node_api.process_transaction_records(records=action_scope.side_effects.transactions)
-        await full_node_api.wait_for_wallets_synced(wallet_nodes=[wallet_node, wallet_node_2])
-
-        expected_confirmed_balance -= fee
-
-        await full_node_api.wait_for_wallet_synced(wallet_node=wallet_node, timeout=15)
 
-        assert await did_wallet_1.get_confirmed_balance() == did_amount
-        assert await did_wallet_1.get_unconfirmed_balance() == did_amount
+        await wallet_environments.process_pending_states(
+            [
+                WalletStateTransition(
+                    pre_block_balance_updates={
+                        "xch": {
+                            "unconfirmed_wallet_balance": -fee,
+                            "set_remainder": True,
+                        },
+                        "did": {
+                            "unconfirmed_wallet_balance": 0,
+                            "set_remainder": True,
+                        },
+                    },
+                    post_block_balance_updates={
+                        "xch": {
+                            "confirmed_wallet_balance": -fee,
+                            "set_remainder": True,
+                        },
+                        "did": {
+                            "confirmed_wallet_balance": 0,
+                            "set_remainder": True,
+                        },
+                    },
+                ),
+            ]
+        )
 
         assert get_parent_num(did_wallet_1) == parent_num + 2
+        assert did_wallet_1.did_info.current_inner is not None  # mypy
         assert puzhash != did_wallet_1.did_info.current_inner.get_tree_hash()
-        assert await wallet.get_confirmed_balance() == expected_confirmed_balance
-        assert await wallet.get_unconfirmed_balance() == expected_confirmed_balance
-
         assert did_wallet_1.did_info.metadata.find("Twitter") > 0
 
-    @pytest.mark.parametrize(
-        "trusted",
-        [True, False],
-    )
+    @pytest.mark.limit_consensus_modes(allowed=[ConsensusMode.PLAIN], reason="irrelevant")
+    @pytest.mark.parametrize("wallet_environments", [{"num_environments": 1, "blocks_needed": [1]}], indirect=True)
     @pytest.mark.anyio
-    async def test_did_sign_message(self, self_hostname, two_wallet_nodes, trusted):
+    async def test_did_sign_message(self, wallet_environments: WalletTestFramework):
+        env = wallet_environments.environments[0]
+        wallet_node = env.node
+        wallet = env.xch_wallet
+        api_0 = env.rpc_api
+
+        env.wallet_aliases = {
+            "xch": 1,
+            "did": 2,
+        }
         fee = uint64(1000)
-        full_nodes, wallets, _ = two_wallet_nodes
-        full_node_api = full_nodes[0]
-        server_1 = full_node_api.server
-        wallet_node, server_2 = wallets[0]
-        wallet_node_2, server_3 = wallets[1]
-        wallet = wallet_node.wallet_state_manager.main_wallet
-        api_0 = WalletRpcApi(wallet_node)
         ph = await wallet.get_new_puzzlehash()
 
-        if trusted:
-            wallet_node.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.config["trusted_peers"] = {}
-            wallet_node_2.config["trusted_peers"] = {}
-
-        await server_2.start_client(PeerInfo(self_hostname, server_1.get_port()), None)
-        await server_3.start_client(PeerInfo(self_hostname, server_1.get_port()), None)
-        await full_node_api.farm_blocks_to_wallet(1, wallet)
-
-        async with wallet.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
+        async with wallet.wallet_state_manager.new_action_scope(
+            wallet_environments.tx_config, push=True
+        ) as action_scope:
             did_wallet_1: DIDWallet = await DIDWallet.create_new_did_wallet(
                 wallet_node.wallet_state_manager,
                 wallet,
                 uint64(101),
                 action_scope,
-                [bytes(ph)],
+                [ph],
                 uint64(1),
                 {"Twitter": "Test", "GitHub": "测试"},
                 fee=fee,
             )
         assert did_wallet_1.get_name() == "Profile 1"
-        await full_node_api.process_transaction_records(records=action_scope.side_effects.transactions)
-        await full_node_api.wait_for_wallets_synced(wallet_nodes=[wallet_node, wallet_node_2])
-        await time_out_assert(15, did_wallet_1.get_confirmed_balance, 101)
+
+        await wallet_environments.process_pending_states(
+            [
+                WalletStateTransition(
+                    pre_block_balance_updates={
+                        "xch": {
+                            "set_remainder": True,
+                        },
+                        "did": {
+                            "init": True,
+                            "unconfirmed_wallet_balance": 101,
+                            "pending_change": 101,
+                            "pending_coin_removal_count": 1,
+                        },
+                    },
+                    post_block_balance_updates={
+                        "xch": {
+                            "set_remainder": True,
+                        },
+                        "did": {
+                            "confirmed_wallet_balance": 101,
+                            "spendable_balance": 101,
+                            "max_send_amount": 101,
+                            "unspent_coin_count": 1,
+                            "pending_change": -101,
+                            "pending_coin_removal_count": -1,
+                        },
+                    },
+                ),
+            ]
+        )
         # Test general string
+        assert did_wallet_1.did_info.origin_coin is not None  # mypy
         message = "Hello World"
+        assert did_wallet_1.did_info.origin_coin is not None
         response = await api_0.sign_message_by_id(
             {
                 "id": encode_puzzle_hash(did_wallet_1.did_info.origin_coin.name(), AddressType.DID.value),
@@ -1213,7 +2003,7 @@ class TestDIDWallet:
                 "is_hex": True,
             }
         )
-        puzzle: Program = Program.to((CHIP_0002_SIGN_MESSAGE_PREFIX, bytes.fromhex(message)))
+        puzzle = Program.to((CHIP_0002_SIGN_MESSAGE_PREFIX, bytes.fromhex(message)))
 
         assert AugSchemeMPL.verify(
             G1Element.from_bytes(bytes.fromhex(response["pubkey"])),
@@ -1223,6 +2013,7 @@ class TestDIDWallet:
 
         # Test BLS sign string
         message = "Hello World"
+        assert did_wallet_1.did_info.origin_coin is not None
         response = await api_0.sign_message_by_id(
             {
                 "id": encode_puzzle_hash(did_wallet_1.did_info.origin_coin.name(), AddressType.DID.value),
@@ -1239,6 +2030,7 @@ class TestDIDWallet:
         )
         # Test BLS sign hex
         message = "0123456789ABCDEF"
+        assert did_wallet_1.did_info.origin_coin is not None
         response = await api_0.sign_message_by_id(
             {
                 "id": encode_puzzle_hash(did_wallet_1.did_info.origin_coin.name(), AddressType.DID.value),
@@ -1259,7 +2051,9 @@ class TestDIDWallet:
         [True, False],
     )
     @pytest.mark.anyio
-    async def test_create_did_with_recovery_list(self, self_hostname, two_nodes_two_wallets_with_same_keys, trusted):
+    async def test_create_did_with_recovery_list(
+        self, self_hostname: str, two_nodes_two_wallets_with_same_keys: OldSimulatorsAndWallets, trusted: bool
+    ) -> None:
         """
         A DID is created on-chain in client0, causing a DID Wallet to be created in client1, which shares the same key.
         This can happen if someone uses the same key on multiple computers, or is syncing a wallet from scratch.
@@ -1303,7 +2097,7 @@ class TestDIDWallet:
 
         # Node 0 sets up a DID Wallet with a backup set, but num_of_backup_ids_needed=0
         # (a malformed solution, but legal for the clvm puzzle)
-        recovery_list = [bytes.fromhex("00" * 32)]
+        recovery_list = [bytes32(bytes.fromhex("00" * 32))]
         async with wallet_0.wallet_state_manager.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
             did_wallet_0: DIDWallet = await DIDWallet.create_new_did_wallet(
                 wallet_node_0.wallet_state_manager,
@@ -1311,7 +2105,7 @@ class TestDIDWallet:
                 uint64(101),
                 action_scope,
                 backups_ids=recovery_list,
-                num_of_backup_ids_needed=0,
+                num_of_backup_ids_needed=uint64(0),
             )
 
         await full_node_api.process_transaction_records(records=action_scope.side_effects.transactions)
@@ -1336,12 +2130,19 @@ class TestDIDWallet:
         #    == json.loads(all_node_1_wallets[1].data)["current_inner"]
         # )
 
+    #  TODO: See Issue CHIA-1544
+    #  This test should be ported to WalletTestFramework once we can replace keys in the wallet node
     @pytest.mark.parametrize(
         "trusted",
         [True, False],
     )
     @pytest.mark.anyio
-    async def test_did_resync(self, self_hostname, two_wallet_nodes, trusted) -> None:
+    async def test_did_resync(
+        self,
+        self_hostname: str,
+        two_wallet_nodes: tuple[list[FullNodeSimulator], list[tuple[WalletNode, ChiaServer]], BlockTools],
+        trusted: bool,
+    ) -> None:
         full_nodes, wallets, _ = two_wallet_nodes
         full_node_api = full_nodes[0]
         full_node_server = full_node_api.full_node.server
@@ -1359,6 +2160,7 @@ class TestDIDWallet:
         else:
             wallet_node_1.config["trusted_peers"] = {}
             wallet_node_2.config["trusted_peers"] = {}
+        assert full_node_server._port is not None
         await wallet_server_1.start_client(PeerInfo(self_hostname, uint16(full_node_server._port)), None)
         await wallet_server_2.start_client(PeerInfo(self_hostname, uint16(full_node_server._port)), None)
         await full_node_api.farm_blocks_to_wallet(1, wallet)
@@ -1408,6 +2210,7 @@ class TestDIDWallet:
         # Start resync
         await wallet_node_1._start_with_fingerprint(fingerprint_1)
         await wallet_node_2._start_with_fingerprint(fingerprint_2)
+        assert full_node_server._port is not None
         await wallet_server_1.start_client(PeerInfo(self_hostname, uint16(full_node_server._port)), None)
         await wallet_server_2.start_client(PeerInfo(self_hostname, uint16(full_node_server._port)), None)
         await full_node_api.farm_new_transaction_block(FarmNewBlockProtocol(bytes32(b"\00" * 32)))