CHIA-619 Annotate test_trades.py (Chia-Network#17447)

Annotate test_trades.py.

chia/_tests/wallet/cat_wallet/test_trades.py

@@ -1,19 +1,20 @@
 from __future__ import annotations
 
 import dataclasses
-from typing import Any, Dict, List, Union
+from typing import Any, Dict, List, Tuple, Union
 
 import pytest
 from chia_rs import G2Element
 
 from chia._tests.conftest import SOFTFORK_HEIGHTS
-from chia._tests.environments.wallet import WalletEnvironment, WalletStateTransition, WalletTestFramework
+from chia._tests.environments.wallet import WalletStateTransition, WalletTestFramework
 from chia._tests.util.time_out_assert import time_out_assert
 from chia._tests.wallet.vc_wallet.test_vc_wallet import mint_cr_cat
 from chia.consensus.cost_calculator import NPCResult
 from chia.consensus.default_constants import DEFAULT_CONSTANTS
 from chia.full_node.bundle_tools import simple_solution_generator
 from chia.full_node.mempool_check_conditions import get_name_puzzle_conditions
+from chia.simulator.full_node_simulator import FullNodeSimulator
 from chia.types.blockchain_format.program import INFINITE_COST, Program
 from chia.types.blockchain_format.sized_bytes import bytes32
 from chia.types.spend_bundle import SpendBundle
@@ -34,6 +35,9 @@
 from chia.wallet.vc_wallet.cr_cat_drivers import ProofsChecker
 from chia.wallet.vc_wallet.cr_cat_wallet import CRCATWallet
 from chia.wallet.vc_wallet.vc_store import VCProofs
+from chia.wallet.wallet_node import WalletNode
+
+OfferSummary = Dict[Union[int, bytes32], int]
 
 
 async def get_trade_and_status(trade_manager: TradeManager, trade: TradeRecord) -> TradeStatus:
@@ -98,10 +102,10 @@ async def test_cat_trades(
     wallet_environments: WalletTestFramework,
     credential_restricted: bool,
     active_softfork_height: uint32,
-):
+) -> None:
     # Setup
-    env_maker: WalletEnvironment = wallet_environments.environments[0]
-    env_taker: WalletEnvironment = wallet_environments.environments[1]
+    env_maker = wallet_environments.environments[0]
+    env_taker = wallet_environments.environments[1]
     wallet_node_maker = env_maker.node
     wallet_node_taker = env_taker.node
     client_maker = env_maker.rpc_client
@@ -136,20 +140,20 @@ async def test_cat_trades(
         }
 
         # Mint some DIDs
-        did_wallet_maker: DIDWallet = await DIDWallet.create_new_did_wallet(
+        did_wallet_maker = await DIDWallet.create_new_did_wallet(
             wallet_node_maker.wallet_state_manager, wallet_maker, uint64(1), wallet_environments.tx_config
         )
-        did_wallet_taker: DIDWallet = await DIDWallet.create_new_did_wallet(
+        did_wallet_taker = await DIDWallet.create_new_did_wallet(
             wallet_node_taker.wallet_state_manager, wallet_taker, uint64(1), wallet_environments.tx_config
         )
         did_id_maker = bytes32.from_hexstr(did_wallet_maker.get_my_DID())
         did_id_taker = bytes32.from_hexstr(did_wallet_taker.get_my_DID())
 
         # Mint some CR-CATs
-        tail_maker: Program = Program.to([3, (1, "maker"), None, None])
-        tail_taker: Program = Program.to([3, (1, "taker"), None, None])
-        proofs_checker_maker: ProofsChecker = ProofsChecker(["foo", "bar"])
-        proofs_checker_taker: ProofsChecker = ProofsChecker(["bar", "zap"])
+        tail_maker = Program.to([3, (1, "maker"), None, None])
+        tail_taker = Program.to([3, (1, "taker"), None, None])
+        proofs_checker_maker = ProofsChecker(["foo", "bar"])
+        proofs_checker_taker = ProofsChecker(["bar", "zap"])
         authorized_providers: List[bytes32] = [did_id_maker, did_id_taker]
         cat_wallet_maker: CATWallet = await CRCATWallet.get_or_create_wallet_for_cat(
             wallet_node_maker.wallet_state_manager,
@@ -251,15 +255,15 @@ async def test_cat_trades(
             ]
         )
 
-        proofs_maker: VCProofs = VCProofs({"foo": "1", "bar": "1", "zap": "1"})
+        proofs_maker = VCProofs({"foo": "1", "bar": "1", "zap": "1"})
         proof_root_maker: bytes32 = proofs_maker.root()
         await client_maker.vc_spend(
             vc_record_maker.vc.launcher_id,
             wallet_environments.tx_config,
             new_proof_hash=proof_root_maker,
         )
 
-        proofs_taker: VCProofs = VCProofs({"foo": "1", "bar": "1", "zap": "1"})
+        proofs_taker = VCProofs({"foo": "1", "bar": "1", "zap": "1"})
         proof_root_taker: bytes32 = proofs_taker.root()
         await client_taker.vc_spend(
             vc_record_taker.vc.launcher_id,
@@ -393,7 +397,6 @@ async def test_cat_trades(
     await env_maker.check_balances()
 
     # Create the trade parameters
-    OfferSummary = Dict[Union[int, bytes32], int]
     chia_for_cat: OfferSummary = {
         wallet_maker.id(): -1,
         bytes32.from_hexstr(new_cat_wallet_maker.get_asset_id()): 2,  # This is the CAT that the taker made
@@ -659,7 +662,7 @@ async def test_cat_trades(
     await time_out_assert(15, get_trade_and_status, TradeStatus.CONFIRMED, trade_manager_maker, trade_make)
     await time_out_assert(15, get_trade_and_status, TradeStatus.CONFIRMED, trade_manager_taker, trade_take)
 
-    async def assert_trade_tx_number(wallet_node, trade_id, number):
+    async def assert_trade_tx_number(wallet_node: WalletNode, trade_id: bytes32, number: int) -> bool:
         txs = await wallet_node.wallet_state_manager.tx_store.get_transactions_by_trade_id(trade_id)
         return len(txs) == number
 
@@ -1580,12 +1583,10 @@ async def assert_trade_tx_number(wallet_node, trade_id, number):
 
 @pytest.mark.parametrize("trusted", [True, False])
 @pytest.mark.anyio
-async def test_trade_cancellation(wallets_prefarm):
-    (
-        [wallet_node_maker, maker_funds],
-        [wallet_node_taker, taker_funds],
-        full_node,
-    ) = wallets_prefarm
+async def test_trade_cancellation(
+    wallets_prefarm: Tuple[Tuple[WalletNode, int], Tuple[WalletNode, int], FullNodeSimulator]
+) -> None:
+    (wallet_node_maker, maker_funds), (wallet_node_taker, taker_funds), full_node = wallets_prefarm
     wallet_maker = wallet_node_maker.wallet_state_manager.main_wallet
     wallet_taker = wallet_node_taker.wallet_state_manager.main_wallet
 
@@ -1607,12 +1608,12 @@ async def test_trade_cancellation(wallets_prefarm):
     maker_funds -= xch_to_cat_amount
     await time_out_assert(15, wallet_maker.get_confirmed_balance, maker_funds)
 
-    cat_for_chia = {
+    cat_for_chia: OfferSummary = {
         wallet_maker.id(): 1,
         cat_wallet_maker.id(): -2,
     }
 
-    chia_for_cat = {
+    chia_for_cat: OfferSummary = {
         wallet_maker.id(): -3,
         cat_wallet_maker.id(): 4,
     }
@@ -1662,6 +1663,7 @@ async def test_trade_cancellation(wallets_prefarm):
     )
     txs = await wallet_maker.wallet_state_manager.add_pending_transactions(txs)
     await time_out_assert(15, get_trade_and_status, TradeStatus.PENDING_CANCEL, trade_manager_maker, trade_make)
+    assert txs is not None
     await full_node.process_transaction_records(records=txs)
 
     sum_of_outgoing = uint64(0)
@@ -1703,6 +1705,7 @@ async def test_trade_cancellation(wallets_prefarm):
     )
     txs = await wallet_maker.wallet_state_manager.add_pending_transactions(txs)
     await time_out_assert(15, get_trade_and_status, TradeStatus.PENDING_CANCEL, trade_manager_maker, trade_make)
+    assert txs is not None
     await full_node.process_transaction_records(records=txs)
 
     await time_out_assert(15, get_trade_and_status, TradeStatus.CANCELLED, trade_manager_maker, trade_make)
@@ -1755,12 +1758,10 @@ async def test_trade_cancellation(wallets_prefarm):
 
 @pytest.mark.parametrize("trusted", [True, False])
 @pytest.mark.anyio
-async def test_trade_cancellation_balance_check(wallets_prefarm):
-    (
-        [wallet_node_maker, maker_funds],
-        [wallet_node_taker, taker_funds],
-        full_node,
-    ) = wallets_prefarm
+async def test_trade_cancellation_balance_check(
+    wallets_prefarm: Tuple[Tuple[WalletNode, int], Tuple[WalletNode, int], FullNodeSimulator]
+) -> None:
+    (wallet_node_maker, maker_funds), _, full_node = wallets_prefarm
     wallet_maker = wallet_node_maker.wallet_state_manager.main_wallet
 
     xch_to_cat_amount = uint64(100)
@@ -1781,7 +1782,7 @@ async def test_trade_cancellation_balance_check(wallets_prefarm):
     maker_funds -= xch_to_cat_amount
     await time_out_assert(15, wallet_maker.get_confirmed_balance, maker_funds)
 
-    chia_for_cat = {
+    chia_for_cat: OfferSummary = {
         wallet_maker.id(): -(await wallet_maker.get_spendable_balance()),
         cat_wallet_maker.id(): 4,
     }
@@ -1798,20 +1799,22 @@ async def test_trade_cancellation_balance_check(wallets_prefarm):
     )
     txs = await trade_manager_maker.wallet_state_manager.add_pending_transactions(txs)
     await time_out_assert(15, get_trade_and_status, TradeStatus.PENDING_CANCEL, trade_manager_maker, trade_make)
+    assert txs is not None
     await full_node.process_transaction_records(records=txs)
 
     await time_out_assert(15, get_trade_and_status, TradeStatus.CANCELLED, trade_manager_maker, trade_make)
 
 
 @pytest.mark.parametrize("trusted", [True, False])
 @pytest.mark.anyio
-async def test_trade_conflict(three_wallets_prefarm):
-    (
-        [wallet_node_maker, maker_funds],
-        [wallet_node_taker, taker_funds],
-        [wallet_node_trader, trader_funds],
-        full_node,
-    ) = three_wallets_prefarm
+async def test_trade_conflict(
+    three_wallets_prefarm: Tuple[
+        Tuple[WalletNode, int], Tuple[WalletNode, int], Tuple[WalletNode, int], FullNodeSimulator
+    ]
+) -> None:
+    ((wallet_node_maker, maker_funds), (wallet_node_taker, _), (wallet_node_trader, _), full_node) = (
+        three_wallets_prefarm
+    )
     wallet_maker = wallet_node_maker.wallet_state_manager.main_wallet
     xch_to_cat_amount = uint64(100)
 
@@ -1831,7 +1834,7 @@ async def test_trade_conflict(three_wallets_prefarm):
     maker_funds -= xch_to_cat_amount
     await time_out_assert(15, wallet_maker.get_confirmed_balance, maker_funds)
 
-    chia_for_cat = {
+    chia_for_cat: OfferSummary = {
         wallet_maker.id(): 1000,
         cat_wallet_maker.id(): -4,
     }
@@ -1872,12 +1875,10 @@ async def test_trade_conflict(three_wallets_prefarm):
 
 @pytest.mark.parametrize("trusted", [True, False])
 @pytest.mark.anyio
-async def test_trade_bad_spend(wallets_prefarm):
-    (
-        [wallet_node_maker, maker_funds],
-        [wallet_node_taker, taker_funds],
-        full_node,
-    ) = wallets_prefarm
+async def test_trade_bad_spend(
+    wallets_prefarm: Tuple[Tuple[WalletNode, int], Tuple[WalletNode, int], FullNodeSimulator]
+) -> None:
+    (wallet_node_maker, maker_funds), (wallet_node_taker, _), full_node = wallets_prefarm
     wallet_maker = wallet_node_maker.wallet_state_manager.main_wallet
     xch_to_cat_amount = uint64(100)
 
@@ -1897,7 +1898,7 @@ async def test_trade_bad_spend(wallets_prefarm):
     maker_funds -= xch_to_cat_amount
     await time_out_assert(15, wallet_maker.get_confirmed_balance, maker_funds)
 
-    chia_for_cat = {
+    chia_for_cat: OfferSummary = {
         wallet_maker.id(): 1000,
         cat_wallet_maker.id(): -4,
     }
@@ -1931,12 +1932,10 @@ def check_wallet_cache_empty() -> bool:
 
 @pytest.mark.parametrize("trusted", [True, False])
 @pytest.mark.anyio
-async def test_trade_high_fee(wallets_prefarm):
-    (
-        [wallet_node_maker, maker_funds],
-        [wallet_node_taker, taker_funds],
-        full_node,
-    ) = wallets_prefarm
+async def test_trade_high_fee(
+    wallets_prefarm: Tuple[Tuple[WalletNode, int], Tuple[WalletNode, int], FullNodeSimulator]
+) -> None:
+    (wallet_node_maker, maker_funds), (wallet_node_taker, _), full_node = wallets_prefarm
     wallet_maker = wallet_node_maker.wallet_state_manager.main_wallet
     xch_to_cat_amount = uint64(100)
 
@@ -1956,7 +1955,7 @@ async def test_trade_high_fee(wallets_prefarm):
     maker_funds -= xch_to_cat_amount
     await time_out_assert(15, wallet_maker.get_confirmed_balance, maker_funds)
 
-    chia_for_cat = {
+    chia_for_cat: OfferSummary = {
         wallet_maker.id(): 1000,
         cat_wallet_maker.id(): -4,
     }
@@ -1984,12 +1983,10 @@ async def test_trade_high_fee(wallets_prefarm):
 
 @pytest.mark.parametrize("trusted", [True, False])
 @pytest.mark.anyio
-async def test_aggregated_trade_state(wallets_prefarm):
-    (
-        [wallet_node_maker, maker_funds],
-        [wallet_node_taker, taker_funds],
-        full_node,
-    ) = wallets_prefarm
+async def test_aggregated_trade_state(
+    wallets_prefarm: Tuple[Tuple[WalletNode, int], Tuple[WalletNode, int], FullNodeSimulator]
+) -> None:
+    (wallet_node_maker, maker_funds), (wallet_node_taker, _), full_node = wallets_prefarm
     wallet_maker = wallet_node_maker.wallet_state_manager.main_wallet
     xch_to_cat_amount = uint64(100)
 
@@ -2009,11 +2006,11 @@ async def test_aggregated_trade_state(wallets_prefarm):
     maker_funds -= xch_to_cat_amount
     await time_out_assert(15, wallet_maker.get_confirmed_balance, maker_funds)
 
-    chia_for_cat = {
+    chia_for_cat: OfferSummary = {
         wallet_maker.id(): 2,
         cat_wallet_maker.id(): -2,
     }
-    cat_for_chia = {
+    cat_for_chia: OfferSummary = {
         wallet_maker.id(): -1,
         cat_wallet_maker.id(): 1,
     }

mypy-exclusions.txt

@@ -93,7 +93,6 @@ chia._tests.util.test_full_block_utils
 chia._tests.util.test_misc
 chia._tests.util.test_network
 chia._tests.util.time_out_assert
-chia._tests.wallet.cat_wallet.test_trades
 chia._tests.wallet.did_wallet.test_did
 chia._tests.wallet.rpc.test_wallet_rpc
 chia._tests.wallet.test_taproot