chia-blockchain 2.5.7rc4__py3-none-any.whl → 2.6.0rc2__py3-none-any.whl

chia/_tests/wallet/sync/test_wallet_sync.py

@@ -4,9 +4,8 @@ import asyncio
 import contextlib
 import functools
 import logging
-from collections.abc import Awaitable
+from collections.abc import Awaitable, Callable
 from dataclasses import dataclass
-from typing import Callable, Optional
 from unittest.mock import MagicMock
 
 import pytest
@@ -26,7 +25,12 @@ from chiabip158 import PyBIP158
 from colorlog import getLogger
 
 from chia._tests.conftest import ConsensusMode
-from chia._tests.connection_utils import connect_and_get_peer, disconnect_all, disconnect_all_and_reconnect
+from chia._tests.connection_utils import (
+    add_dummy_connection,
+    connect_and_get_peer,
+    disconnect_all,
+    disconnect_all_and_reconnect,
+)
 from chia._tests.util.blockchain_mock import BlockchainMock
 from chia._tests.util.misc import patch_request_handler, wallet_height_at_least
 from chia._tests.util.setup_nodes import OldSimulatorsAndWallets
@@ -111,7 +115,9 @@ async def test_request_block_headers(
     assert [x.transactions_filter for x in bh] == [b"\x00"] * 6
 
     num_blocks = 20
-    new_blocks = bt.get_consecutive_blocks(num_blocks, block_list_input=default_400_blocks, pool_reward_puzzle_hash=ph)
+    new_blocks = bt.get_consecutive_blocks(
+        num_blocks, block_list_input=default_400_blocks, farmer_reward_puzzle_hash=ph
+    )
     await add_blocks_in_batches(new_blocks, full_node_api.full_node)
     msg = await full_node_api.request_block_headers(wallet_protocol.RequestBlockHeaders(uint32(110), uint32(115), True))
     assert msg is not None
@@ -626,7 +632,7 @@ async def test_request_additions_errors(simulator_and_wallet: OldSimulatorsAndWa
 
     await full_node_api.wait_for_wallet_synced(wallet_node=wallet_node, timeout=20)
 
-    last_block: Optional[BlockRecord] = full_node_api.full_node.blockchain.get_peak()
+    last_block: BlockRecord | None = full_node_api.full_node.blockchain.get_peak()
     assert last_block is not None
 
     # Invalid height
@@ -675,7 +681,7 @@ async def test_request_additions_success(simulator_and_wallet: OldSimulatorsAndW
         ph = await action_scope.get_puzzle_hash(wallet.wallet_state_manager)
 
     full_node_api = full_nodes[0]
-    await wallet_server.start_client(PeerInfo(self_hostname, full_node_api.full_node.server.get_port()), None)
+    wallet_peer = await connect_and_get_peer(full_node_api.full_node.server, wallet_server, self_hostname)
 
     for _ in range(2):
         await full_node_api.farm_new_transaction_block(FarmNewBlockProtocol(ph))
@@ -697,7 +703,7 @@ async def test_request_additions_success(simulator_and_wallet: OldSimulatorsAndW
         )
     [tx] = action_scope.side_effects.transactions
     assert tx.spend_bundle is not None
-    await full_node_api.send_transaction(SendTransaction(tx.spend_bundle))
+    await full_node_api.send_transaction(SendTransaction(tx.spend_bundle), wallet_peer)
     await full_node_api.wait_transaction_records_entered_mempool([tx])
     await full_node_api.farm_new_transaction_block(FarmNewBlockProtocol(ph))
 
@@ -917,7 +923,9 @@ async def test_dusted_wallet(
         )
     [tx] = action_scope.side_effects.transactions
     assert tx.spend_bundle is not None
-    await full_node_api.send_transaction(SendTransaction(tx.spend_bundle))
+    _, dummy_node_id = await add_dummy_connection(full_node_api.server, self_hostname, 12312)
+    dummy_peer = full_node_api.server.all_connections[dummy_node_id]
+    await full_node_api.send_transaction(SendTransaction(tx.spend_bundle), dummy_peer)
     await full_node_api.wait_transaction_records_entered_mempool([tx])
     await full_node_api.wait_for_wallets_synced(wallet_nodes=[farm_wallet_node, dust_wallet_node], timeout=20)
     await full_node_api.farm_new_transaction_block(FarmNewBlockProtocol(ph))
@@ -984,7 +992,7 @@ async def test_dusted_wallet(
                 )
             [tx] = action_scope.side_effects.transactions
             assert tx.spend_bundle is not None
-            await full_node_api.send_transaction(SendTransaction(tx.spend_bundle))
+            await full_node_api.send_transaction(SendTransaction(tx.spend_bundle), dummy_peer)
 
             # advance the chain and sync both wallets
             await full_node_api.wait_transaction_records_entered_mempool([tx])
@@ -1007,7 +1015,7 @@ async def test_dusted_wallet(
             )
         [tx] = action_scope.side_effects.transactions
         assert tx.spend_bundle is not None
-        await full_node_api.send_transaction(SendTransaction(tx.spend_bundle))
+        await full_node_api.send_transaction(SendTransaction(tx.spend_bundle), dummy_peer)
 
         # advance the chain and sync both wallets
         await full_node_api.wait_transaction_records_entered_mempool([tx])
@@ -1060,7 +1068,7 @@ async def test_dusted_wallet(
         )
     [tx] = action_scope.side_effects.transactions
     assert tx.spend_bundle is not None
-    await full_node_api.send_transaction(SendTransaction(tx.spend_bundle))
+    await full_node_api.send_transaction(SendTransaction(tx.spend_bundle), dummy_peer)
 
     # advance the chain and sync both wallets
     await full_node_api.wait_transaction_records_entered_mempool([tx])
@@ -1104,7 +1112,7 @@ async def test_dusted_wallet(
         )
     [tx] = action_scope.side_effects.transactions
     assert tx.spend_bundle is not None
-    await full_node_api.send_transaction(SendTransaction(tx.spend_bundle))
+    await full_node_api.send_transaction(SendTransaction(tx.spend_bundle), dummy_peer)
 
     # advance the chain and sync both wallets
     await full_node_api.wait_transaction_records_entered_mempool([tx])
@@ -1168,7 +1176,7 @@ async def test_dusted_wallet(
         )
     [tx] = action_scope.side_effects.transactions
     assert tx.spend_bundle is not None
-    await full_node_api.send_transaction(SendTransaction(tx.spend_bundle))
+    await full_node_api.send_transaction(SendTransaction(tx.spend_bundle), dummy_peer)
 
     # advance the chain and sync both wallets
     await full_node_api.wait_transaction_records_entered_mempool([tx])
@@ -1209,7 +1217,7 @@ async def test_dusted_wallet(
         )
     [tx] = action_scope.side_effects.transactions
     assert tx.spend_bundle is not None
-    await full_node_api.send_transaction(SendTransaction(tx.spend_bundle))
+    await full_node_api.send_transaction(SendTransaction(tx.spend_bundle), dummy_peer)
 
     # advance the chain and sync both wallets
     await full_node_api.wait_transaction_records_entered_mempool([tx])
@@ -1258,7 +1266,7 @@ async def test_dusted_wallet(
                 )
             [tx] = action_scope.side_effects.transactions
             assert tx.spend_bundle is not None
-            await full_node_api.send_transaction(SendTransaction(tx.spend_bundle))
+            await full_node_api.send_transaction(SendTransaction(tx.spend_bundle), dummy_peer)
             await full_node_api.wait_transaction_records_entered_mempool([tx])
             await full_node_api.wait_for_wallets_synced(wallet_nodes=[farm_wallet_node, dust_wallet_node], timeout=20)
             await full_node_api.farm_new_transaction_block(FarmNewBlockProtocol(ph))
@@ -1277,7 +1285,7 @@ async def test_dusted_wallet(
         )
     [tx] = action_scope.side_effects.transactions
     assert tx.spend_bundle is not None
-    await full_node_api.send_transaction(SendTransaction(tx.spend_bundle))
+    await full_node_api.send_transaction(SendTransaction(tx.spend_bundle), dummy_peer)
 
     # advance the chain and sync both wallets
     await full_node_api.wait_transaction_records_entered_mempool([tx])
@@ -1312,7 +1320,7 @@ async def test_dusted_wallet(
         )
     [tx] = action_scope.side_effects.transactions
     assert tx.spend_bundle is not None
-    await full_node_api.send_transaction(SendTransaction(tx.spend_bundle))
+    await full_node_api.send_transaction(SendTransaction(tx.spend_bundle), dummy_peer)
 
     # advance the chain and sync both wallets
     await full_node_api.wait_transaction_records_entered_mempool([tx])
@@ -1431,10 +1439,10 @@ async def test_retry_store(
 
     def flaky_get_coin_state(
         flakiness_info: FlakinessInfo,
-        func: Callable[[list[bytes32], WSChiaConnection, Optional[uint32]], Awaitable[list[CoinState]]],
-    ) -> Callable[[list[bytes32], WSChiaConnection, Optional[uint32]], Awaitable[list[CoinState]]]:
+        func: Callable[[list[bytes32], WSChiaConnection, uint32 | None], Awaitable[list[CoinState]]],
+    ) -> Callable[[list[bytes32], WSChiaConnection, uint32 | None], Awaitable[list[CoinState]]]:
         async def new_func(
-            coin_names: list[bytes32], peer: WSChiaConnection, fork_height: Optional[uint32] = None
+            coin_names: list[bytes32], peer: WSChiaConnection, fork_height: uint32 | None = None
         ) -> list[CoinState]:
             if flakiness_info.coin_state_flaky:
                 flakiness_info.coin_state_flaky = False
@@ -1447,10 +1455,10 @@ async def test_retry_store(
     request_puzzle_solution_failure_tested = False
 
     def flaky_request_puzzle_solution(
-        func: Callable[[FullNodeAPI, wallet_protocol.RequestPuzzleSolution], Awaitable[Optional[Message]]],
-    ) -> Callable[[FullNodeAPI, wallet_protocol.RequestPuzzleSolution], Awaitable[Optional[Message]]]:
+        func: Callable[[FullNodeAPI, wallet_protocol.RequestPuzzleSolution], Awaitable[Message | None]],
+    ) -> Callable[[FullNodeAPI, wallet_protocol.RequestPuzzleSolution], Awaitable[Message | None]]:
         @functools.wraps(func)
-        async def new_func(self: FullNodeAPI, request: wallet_protocol.RequestPuzzleSolution) -> Optional[Message]:
+        async def new_func(self: FullNodeAPI, request: wallet_protocol.RequestPuzzleSolution) -> Message | None:
             nonlocal request_puzzle_solution_failure_tested
             if not request_puzzle_solution_failure_tested:
                 request_puzzle_solution_failure_tested = True
@@ -1464,10 +1472,10 @@ async def test_retry_store(
 
     def flaky_fetch_children(
         flakiness_info: FlakinessInfo,
-        func: Callable[[bytes32, WSChiaConnection, Optional[uint32]], Awaitable[list[CoinState]]],
-    ) -> Callable[[bytes32, WSChiaConnection, Optional[uint32]], Awaitable[list[CoinState]]]:
+        func: Callable[[bytes32, WSChiaConnection, uint32 | None], Awaitable[list[CoinState]]],
+    ) -> Callable[[bytes32, WSChiaConnection, uint32 | None], Awaitable[list[CoinState]]]:
         async def new_func(
-            coin_name: bytes32, peer: WSChiaConnection, fork_height: Optional[uint32] = None
+            coin_name: bytes32, peer: WSChiaConnection, fork_height: uint32 | None = None
         ) -> list[CoinState]:
             if flakiness_info.fetch_children_flaky:
                 flakiness_info.fetch_children_flaky = False
@@ -1490,9 +1498,9 @@ async def test_retry_store(
         return new_func
 
     def flaky_info_for_puzhash(
-        flakiness_info: FlakinessInfo, func: Callable[[bytes32], Awaitable[Optional[WalletIdentifier]]]
-    ) -> Callable[[bytes32], Awaitable[Optional[WalletIdentifier]]]:
-        async def new_func(puzzle_hash: bytes32) -> Optional[WalletIdentifier]:
+        flakiness_info: FlakinessInfo, func: Callable[[bytes32], Awaitable[WalletIdentifier | None]]
+    ) -> Callable[[bytes32], Awaitable[WalletIdentifier | None]]:
+        async def new_func(puzzle_hash: bytes32) -> WalletIdentifier | None:
             if flakiness_info.db_flaky:
                 flakiness_info.db_flaky = False
                 raise AIOSqliteError
@@ -1605,7 +1613,7 @@ async def test_bad_peak_mismatch(
         full_node_protocol.RespondProofOfWeight(wp, wp.recent_chain_data[-1].header_hash),
     )
     with monkeypatch.context() as m:
-        f: asyncio.Future[Optional[Message]] = asyncio.Future()
+        f: asyncio.Future[Message | None] = asyncio.Future()
         f.set_result(wp_msg)
         m.setattr(full_node_api, "request_proof_of_weight", MagicMock(return_value=f))
 
@@ -1614,7 +1622,7 @@ async def test_bad_peak_mismatch(
             ProtocolMessageTypes.respond_block_header,
             wallet_protocol.RespondBlockHeader(wp.recent_chain_data[-1]),
         )
-        f2: asyncio.Future[Optional[Message]] = asyncio.Future()
+        f2: asyncio.Future[Message | None] = asyncio.Future()
         f2.set_result(header_block_msg)
         m.setattr(full_node_api, "request_block_header", MagicMock(return_value=f2))
 
@@ -1638,7 +1646,7 @@ async def test_bad_peak_mismatch(
 async def test_long_sync_untrusted_break(
     setup_two_nodes_and_wallet: OldSimulatorsAndWallets,
     default_1000_blocks: list[FullBlock],
-    default_400_blocks: list[FullBlock],
+    default_1500_blocks: list[FullBlock],
     self_hostname: str,
     caplog: pytest.LogCaptureFixture,
     use_delta_sync: bool,
@@ -1649,27 +1657,16 @@ async def test_long_sync_untrusted_break(
     wallet_node.config["trusted_peers"] = {trusted_full_node_server.node_id.hex(): None}
     wallet_node.config["use_delta_sync"] = use_delta_sync
 
-    sync_canceled = False
-
     async def register_for_ph_updates(
         self: object,
         request: wallet_protocol.RegisterForPhUpdates,
         peer: WSChiaConnection,
     ) -> None:
-        nonlocal sync_canceled
-        # Just sleep a long time here to simulate a long-running untrusted sync
-        try:
-            await asyncio.sleep(120)
-        except Exception:
-            sync_canceled = True
-            raise
+        await peer.wait_until_closed()
 
     def wallet_syncing() -> bool:
         return wallet_node.wallet_state_manager.sync_mode
 
-    def check_sync_canceled() -> bool:
-        return sync_canceled
-
     def synced_to_trusted() -> bool:
         return trusted_full_node_server.node_id in wallet_node.synced_peers
 
@@ -1678,26 +1675,32 @@ async def test_long_sync_untrusted_break(
         untrusted_peers = sum(not wallet_node.is_trusted(peer) for peer in wallet_server.all_connections.values())
         return trusted_peers == 1 and untrusted_peers == 0
 
-    await add_blocks_in_batches(default_400_blocks, trusted_full_node_api.full_node)
-
-    await add_blocks_in_batches(default_1000_blocks[:400], untrusted_full_node_api.full_node)
+    await add_blocks_in_batches(default_1500_blocks[:600], trusted_full_node_api.full_node)
+    await add_blocks_in_batches(default_1000_blocks[:600], untrusted_full_node_api.full_node)
 
-    with patch_request_handler(api=untrusted_full_node_api, handler=register_for_ph_updates):
+    with (
+        patch_request_handler(api=untrusted_full_node_api, handler=register_for_ph_updates),
+        caplog.at_level(logging.INFO),
+    ):
         # Connect to the untrusted peer and wait until the long sync started
-        await wallet_server.start_client(PeerInfo(self_hostname, untrusted_full_node_server.get_port()), None)
+        assert await wallet_server.start_client(PeerInfo(self_hostname, untrusted_full_node_server.get_port()), None)
+        untrusted_conn = wallet_server.all_connections.get(untrusted_full_node_server.node_id)
+        assert untrusted_conn is not None
         await time_out_assert(30, wallet_syncing)
-        with caplog.at_level(logging.INFO):
-            # Connect to the trusted peer and make sure the running untrusted long sync gets interrupted via disconnect
-            await wallet_server.start_client(PeerInfo(self_hostname, trusted_full_node_server.get_port()), None)
-            await time_out_assert(600, wallet_height_at_least, True, wallet_node, len(default_400_blocks) - 1)
-            assert time_out_assert(10, synced_to_trusted)
-            assert untrusted_full_node_server.node_id not in wallet_node.synced_peers
-            assert "Connected to a synced trusted peer, disconnecting from all untrusted nodes." in caplog.text
-
-        # Make sure the sync was interrupted
-        assert time_out_assert(30, check_sync_canceled)
+
+        # Connect to the trusted peer and make sure the running untrusted long sync gets interrupted via disconnect
+        assert await wallet_server.start_client(PeerInfo(self_hostname, trusted_full_node_server.get_port()), None)
+        await time_out_assert(600, wallet_height_at_least, True, wallet_node, 599)
+        await time_out_assert(10, synced_to_trusted)
+        assert untrusted_full_node_server.node_id not in wallet_node.synced_peers
+
+        assert "Connected to a synced trusted peer, disconnecting from all untrusted nodes." in caplog.text
+        assert (
+            f"Connection closed: {untrusted_conn.peer_info.host}, node id: {untrusted_full_node_server.node_id}"
+            in caplog.text
+        )
         # And that we only have a trusted peer left
-        assert time_out_assert(30, only_trusted_peer)
+        await time_out_assert(30, only_trusted_peer)
 
 
 @pytest.mark.anyio

chia/_tests/wallet/test_clvm_streamable.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import dataclasses
-from typing import Optional
 
 import pytest
 from chia_rs.sized_bytes import bytes32
@@ -77,14 +76,14 @@ def test_nested_serialization() -> None:
 @streamable
 @dataclasses.dataclass(frozen=True)
 class Compound(Streamable):
-    optional: Optional[BasicCLVMStreamable]
+    optional: BasicCLVMStreamable | None
     list: list[BasicCLVMStreamable]
 
 
 @clvm_streamable
 @dataclasses.dataclass(frozen=True)
 class CompoundCLVM(Streamable):
-    optional: Optional[BasicCLVMStreamable]
+    optional: BasicCLVMStreamable | None
     list: list[BasicCLVMStreamable]
 
 

chia/_tests/wallet/test_coin_management.py

@@ -66,8 +66,8 @@ def test_list_parsing(id: ValueAndArgs, show_unconfirmed: ValueAndArgs, paginate
         {
             "num_environments": 1,
             "blocks_needed": [3],  # 6 coins to test pagination
-            "reuse_puzhash": True,  # irrelevent
-            "trusted": True,  # irrelevent
+            "reuse_puzhash": True,  # irrelevant
+            "trusted": True,  # irrelevant
         }
     ],
     indirect=True,

chia/_tests/wallet/test_conditions.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Any, Optional, Union
+from typing import Any
 
 import pytest
 from chia_rs.sized_bytes import bytes32
@@ -214,11 +214,9 @@ def test_unknown_condition() -> None:
     ],
 )
 def test_announcement_inversions(
-    drivers: Union[
-        tuple[type[CreateCoinAnnouncement], type[AssertCoinAnnouncement]],
-        tuple[type[CreatePuzzleAnnouncement], type[AssertPuzzleAnnouncement]],
-        tuple[type[CreateAnnouncement], type[AssertAnnouncement]],
-    ],
+    drivers: tuple[type[CreateCoinAnnouncement], type[AssertCoinAnnouncement]]
+    | tuple[type[CreatePuzzleAnnouncement], type[AssertPuzzleAnnouncement]]
+    | tuple[type[CreateAnnouncement], type[AssertAnnouncement]],
 ) -> None:
     create_driver, assert_driver = drivers
     # mypy is not smart enough to understand that this `if` narrows down the potential types it could be
@@ -250,7 +248,7 @@ def test_announcement_inversions(
 class TimelockInfo:
     drivers: list[Condition]
     parsed_info: ConditionValidTimes
-    conditions_after: Optional[list[Condition]] = None
+    conditions_after: list[Condition] | None = None
 
 
 @pytest.mark.parametrize(
@@ -285,9 +283,7 @@ class TimelockInfo:
     ],
 )
 def test_timelock_parsing(timelock_info: TimelockInfo) -> None:
-    assert timelock_info.parsed_info == parse_timelock_info(
-        [UnknownCondition(Program.to(None), []), *timelock_info.drivers]
-    )
+    assert timelock_info.parsed_info == parse_timelock_info([UnknownCondition(Program.NIL, []), *timelock_info.drivers])
     assert timelock_info.parsed_info.to_conditions() == (
         timelock_info.conditions_after if timelock_info.conditions_after is not None else timelock_info.drivers
     )
@@ -349,47 +345,45 @@ def test_timelock_parsing(timelock_info: TimelockInfo) -> None:
 )
 def test_invalid_condition(
     cond: type[
-        Union[
-            AggSigParent,
-            AggSigPuzzle,
-            AggSigAmount,
-            AggSigPuzzleAmount,
-            AggSigParentAmount,
-            AggSigParentPuzzle,
-            AggSigUnsafe,
-            AggSigMe,
-            CreateCoin,
-            ReserveFee,
-            AssertCoinAnnouncement,
-            CreateCoinAnnouncement,
-            AssertPuzzleAnnouncement,
-            CreatePuzzleAnnouncement,
-            AssertConcurrentSpend,
-            AssertConcurrentPuzzle,
-            AssertMyCoinID,
-            AssertMyParentID,
-            AssertMyPuzzleHash,
-            AssertMyAmount,
-            AssertMyBirthSeconds,
-            AssertMyBirthHeight,
-            AssertSecondsRelative,
-            AssertSecondsAbsolute,
-            AssertHeightRelative,
-            AssertHeightAbsolute,
-            AssertBeforeSecondsRelative,
-            AssertBeforeSecondsAbsolute,
-            AssertBeforeHeightRelative,
-            AssertBeforeHeightAbsolute,
-            Softfork,
-            Remark,
-            UnknownCondition,
-            AggSig,
-            CreateAnnouncement,
-            AssertAnnouncement,
-            Timelock,
-            SendMessage,
-            ReceiveMessage,
-        ]
+        AggSigParent
+        | AggSigPuzzle
+        | AggSigAmount
+        | AggSigPuzzleAmount
+        | AggSigParentAmount
+        | AggSigParentPuzzle
+        | AggSigUnsafe
+        | AggSigMe
+        | CreateCoin
+        | ReserveFee
+        | AssertCoinAnnouncement
+        | CreateCoinAnnouncement
+        | AssertPuzzleAnnouncement
+        | CreatePuzzleAnnouncement
+        | AssertConcurrentSpend
+        | AssertConcurrentPuzzle
+        | AssertMyCoinID
+        | AssertMyParentID
+        | AssertMyPuzzleHash
+        | AssertMyAmount
+        | AssertMyBirthSeconds
+        | AssertMyBirthHeight
+        | AssertSecondsRelative
+        | AssertSecondsAbsolute
+        | AssertHeightRelative
+        | AssertHeightAbsolute
+        | AssertBeforeSecondsRelative
+        | AssertBeforeSecondsAbsolute
+        | AssertBeforeHeightRelative
+        | AssertBeforeHeightAbsolute
+        | Softfork
+        | Remark
+        | UnknownCondition
+        | AggSig
+        | CreateAnnouncement
+        | AssertAnnouncement
+        | Timelock
+        | SendMessage
+        | ReceiveMessage
     ],
     prg: bytes,
 ) -> None:

chia/_tests/wallet/test_debug_spend_bundle.py

@@ -50,7 +50,7 @@ def test_debug_spend_bundle(capsys: pytest.CaptureFixture[str]) -> None:
                 make_spend(
                     coin_bad_reveal,
                     ACS,
-                    Program.to(None),
+                    Program.NIL,
                 ),
                 make_spend(
                     coin,
@@ -60,7 +60,7 @@ def test_debug_spend_bundle(capsys: pytest.CaptureFixture[str]) -> None:
                 make_spend(
                     child_coin,
                     ACS,
-                    Program.to(None),
+                    Program.NIL,
                 ),
             ],
             sig,

chia/_tests/wallet/test_new_wallet_protocol.py

@@ -5,10 +5,9 @@ from collections import OrderedDict
 from collections.abc import AsyncGenerator
 from dataclasses import dataclass
 from random import Random
-from typing import Optional
 
 import pytest
-from chia_rs import AugSchemeMPL, Coin, CoinSpend, CoinState, Program, SpendBundle
+from chia_rs import AugSchemeMPL, Coin, CoinRecord, CoinSpend, CoinState, Program, SpendBundle
 from chia_rs.sized_bytes import bytes32
 from chia_rs.sized_ints import uint8, uint16, uint32, uint64
 
@@ -26,7 +25,6 @@ from chia.simulator import simulator_protocol
 from chia.simulator.block_tools import BlockTools
 from chia.simulator.full_node_simulator import FullNodeSimulator
 from chia.simulator.start_simulator import SimulatorFullNodeService
-from chia.types.coin_record import CoinRecord
 from chia.types.mempool_inclusion_status import MempoolInclusionStatus
 from chia.util.hash import std_hash
 from chia.wallet.wallet_service import WalletService
@@ -621,14 +619,14 @@ async def test_request_puzzle_state_limit(one_node: OneNode, self_hostname: str)
 class PuzzleStateData:
     coin_states: list[CoinState]
     end_of_batch: bool
-    previous_height: Optional[uint32]
+    previous_height: uint32 | None
     header_hash: bytes32
 
 
 async def sync_puzzle_hashes(
     puzzle_hashes: list[bytes32],
     *,
-    initial_previous_height: Optional[uint32],
+    initial_previous_height: uint32 | None,
     initial_header_hash: bytes32,
     filters: wallet_protocol.CoinStateFilters,
     subscribe_when_finished: bool = False,
@@ -835,7 +833,7 @@ async def raw_mpu_setup(one_node: OneNode, self_hostname: str, no_capability: bo
     for coin, hint in new_coins:
         solution = Program.to([[]])
         bundle = SpendBundle([CoinSpend(coin, puzzle, solution)], AugSchemeMPL.aggregate([]))
-        tx_resp = await simulator.send_transaction(wallet_protocol.SendTransaction(bundle))
+        tx_resp = await simulator.send_transaction(wallet_protocol.SendTransaction(bundle), peer)
         assert tx_resp is not None
 
         ack = wallet_protocol.TransactionAck.from_bytes(tx_resp.data)
@@ -893,11 +891,13 @@ async def subscribe_puzzle(
     assert len(response.coin_states) == existing_coin_states
 
 
-async def spend_coin(simulator: FullNodeSimulator, coin: Coin, solution: Optional[Program] = None) -> bytes32:
+async def spend_coin(
+    simulator: FullNodeSimulator, peer: WSChiaConnection, coin: Coin, solution: Program | None = None
+) -> bytes32:
     bundle = SpendBundle(
         [CoinSpend(coin, IDENTITY_PUZZLE, Program.to([]) if solution is None else solution)], AugSchemeMPL.aggregate([])
     )
-    tx_resp = await simulator.send_transaction(wallet_protocol.SendTransaction(bundle))
+    tx_resp = await simulator.send_transaction(wallet_protocol.SendTransaction(bundle), peer)
     assert tx_resp is not None
 
     ack = wallet_protocol.TransactionAck.from_bytes(tx_resp.data)
@@ -917,7 +917,7 @@ async def test_spent_coin_id_mempool_update(mpu_setup: Mpu) -> None:
     # Make a coin and spend it
     coin, _ = await make_coin(simulator.full_node)
     await subscribe_coin(simulator, coin.name(), peer)
-    transaction_id = await spend_coin(simulator, coin)
+    transaction_id = await spend_coin(simulator, peer, coin)
 
     # We should have gotten a mempool update for this transaction
     await assert_mempool_added(queue, {transaction_id})
@@ -944,7 +944,7 @@ async def test_spent_puzzle_hash_mempool_update(mpu_setup: Mpu) -> None:
     # Make a coin and spend it
     coin, _ = await make_coin(simulator.full_node)
     await subscribe_puzzle(simulator, coin.puzzle_hash, peer)
-    transaction_id = await spend_coin(simulator, coin)
+    transaction_id = await spend_coin(simulator, peer, coin)
 
     # We should have gotten a mempool update for this transaction
     await assert_mempool_added(queue, {transaction_id})
@@ -971,7 +971,7 @@ async def test_spent_hint_mempool_update(mpu_setup: Mpu) -> None:
     # Make a coin and spend it
     coin, hint = await make_coin(simulator.full_node)
     await subscribe_puzzle(simulator, hint, peer)
-    transaction_id = await spend_coin(simulator, coin)
+    transaction_id = await spend_coin(simulator, peer, coin)
 
     # We should have gotten a mempool update for this transaction
     await assert_mempool_added(queue, {transaction_id})
@@ -1000,7 +1000,7 @@ async def test_created_coin_id_mempool_update(mpu_setup: Mpu) -> None:
     child_coin = Coin(coin.name(), std_hash(b"new puzzle hash"), coin.amount)
     await subscribe_coin(simulator, child_coin.name(), peer, existing_coin_states=0)
     transaction_id = await spend_coin(
-        simulator, coin, solution=Program.to([[51, child_coin.puzzle_hash, child_coin.amount]])
+        simulator, peer, coin, solution=Program.to([[51, child_coin.puzzle_hash, child_coin.amount]])
     )
 
     # We should have gotten a mempool update for this transaction
@@ -1030,7 +1030,7 @@ async def test_created_puzzle_hash_mempool_update(mpu_setup: Mpu) -> None:
     child_coin = Coin(coin.name(), std_hash(b"new puzzle hash"), coin.amount)
     await subscribe_puzzle(simulator, child_coin.puzzle_hash, peer, existing_coin_states=0)
     transaction_id = await spend_coin(
-        simulator, coin, solution=Program.to([[51, child_coin.puzzle_hash, child_coin.amount]])
+        simulator, peer, coin, solution=Program.to([[51, child_coin.puzzle_hash, child_coin.amount]])
     )
 
     # We should have gotten a mempool update for this transaction
@@ -1061,7 +1061,7 @@ async def test_created_hint_mempool_update(mpu_setup: Mpu) -> None:
     hint = std_hash(b"new hint")
     await subscribe_puzzle(simulator, hint, peer, existing_coin_states=0)
     transaction_id = await spend_coin(
-        simulator, coin, solution=Program.to([[51, child_coin.puzzle_hash, child_coin.amount, [hint]]])
+        simulator, peer, coin, solution=Program.to([[51, child_coin.puzzle_hash, child_coin.amount, [hint]]])
     )
 
     # We should have gotten a mempool update for this transaction
@@ -1089,7 +1089,7 @@ async def test_missing_capability_coin_id(mpu_setup_no_capability: Mpu) -> None:
     # Make a coin and spend it
     coin, _ = await make_coin(simulator.full_node)
     await subscribe_coin(simulator, coin.name(), peer)
-    transaction_id = await spend_coin(simulator, coin)
+    transaction_id = await spend_coin(simulator, peer, coin)
 
     # There is no mempool update for this transaction since the peer doesn't have the capability
     assert queue.empty()
@@ -1114,7 +1114,7 @@ async def test_missing_capability_puzzle_hash(mpu_setup_no_capability: Mpu) -> N
     # Make a coin and spend it
     coin, _ = await make_coin(simulator.full_node)
     await subscribe_puzzle(simulator, coin.puzzle_hash, peer)
-    transaction_id = await spend_coin(simulator, coin)
+    transaction_id = await spend_coin(simulator, peer, coin)
 
     # There is no mempool update for this transaction since the peer doesn't have the capability
     assert queue.empty()
@@ -1139,7 +1139,7 @@ async def test_missing_capability_hint(mpu_setup_no_capability: Mpu) -> None:
     # Make a coin and spend it
     coin, hint = await make_coin(simulator.full_node)
     await subscribe_puzzle(simulator, hint, peer)
-    transaction_id = await spend_coin(simulator, coin)
+    transaction_id = await spend_coin(simulator, peer, coin)
 
     # There is no mempool update for this transaction since the peer doesn't have the capability
     assert queue.empty()