From 064d34f1844aec94650b2ac1d698c7704f501bf9 Mon Sep 17 00:00:00 2001 From: Matt Hauff Date: Fri, 12 Jul 2024 12:29:09 -0700 Subject: [PATCH] [CHIA-738] Add a better clawback auto claim test and fix related issue (#18141) * Add the concept of 'action scopes' * Add `WalletActionScope` * Fix CATWallet pending_change calculation * Add the concept of 'action scopes' * pylint and test coverage * add try/finally * add try/except * Undo giving a variable a name * Fix CRCAT test * Fix trade tests * Fix cat test * Add auto claim test coverage --- chia/_tests/wallet/test_wallet.py | 115 +++++++++++------------------- chia/wallet/wallet_node.py | 12 ++-- 2 files changed, 48 insertions(+), 79 deletions(-) diff --git a/chia/_tests/wallet/test_wallet.py b/chia/_tests/wallet/test_wallet.py index 9ce2ba9cbd03..82290777c82d 100644 --- a/chia/_tests/wallet/test_wallet.py +++ b/chia/_tests/wallet/test_wallet.py @@ -21,7 +21,7 @@ from chia.types.spend_bundle import estimate_fees from chia.util.bech32m import encode_puzzle_hash from chia.util.errors import Err -from chia.util.ints import uint32, uint64 +from chia.util.ints import uint16, uint32, uint64 from chia.wallet.conditions import ConditionValidTimes from chia.wallet.derive_keys import master_sk_to_wallet_sk from chia.wallet.payment import Payment @@ -163,12 +163,15 @@ async def test_wallet_reuse_address(self, wallet_environments: WalletTestFramewo @pytest.mark.parametrize( "wallet_environments", - [{"num_environments": 2, "blocks_needed": [1, 1], "reuse_puzhash": True}], + [{"num_environments": 2, "blocks_needed": [2, 1], "reuse_puzhash": True}], indirect=True, ) + @pytest.mark.parametrize("number_of_coins", [1, 3]) @pytest.mark.limit_consensus_modes(reason="irrelevant") @pytest.mark.anyio - async def test_wallet_clawback_claim_auto(self, wallet_environments: WalletTestFramework) -> None: + async def test_wallet_clawback_claim_auto( + self, wallet_environments: WalletTestFramework, number_of_coins: int + ) -> None: env = wallet_environments.environments[0] env_1 = wallet_environments.environments[1] wallet = env.xch_wallet @@ -180,34 +183,35 @@ async def test_wallet_clawback_claim_auto(self, wallet_environments: WalletTestF normal_puzhash = await wallet_1.get_new_puzzlehash() # Transfer to normal wallet - [tx1] = await wallet.generate_signed_transaction( - uint64(tx_amount), - normal_puzhash, - DEFAULT_TX_CONFIG, - uint64(0), - puzzle_decorator_override=[{"decorator": "CLAWBACK", "clawback_timelock": 10}], - ) - [tx1] = await wallet.wallet_state_manager.add_pending_transactions([tx1]) + for _ in range(0, number_of_coins): + [tx1] = await wallet.generate_signed_transaction( + uint64(tx_amount), + normal_puzhash, + DEFAULT_TX_CONFIG, + uint64(0), + puzzle_decorator_override=[{"decorator": "CLAWBACK", "clawback_timelock": 10}], + ) + [tx1] = await wallet.wallet_state_manager.add_pending_transactions([tx1]) await wallet_environments.process_pending_states( [ WalletStateTransition( pre_block_balance_updates={ 1: { - "unconfirmed_wallet_balance": -1 * tx_amount, - "<=#spendable_balance": -1 * tx_amount, - "<=#max_send_amount": -1 * tx_amount, + "unconfirmed_wallet_balance": -1 * tx_amount * number_of_coins, + "<=#spendable_balance": -1 * tx_amount * number_of_coins, + "<=#max_send_amount": -1 * tx_amount * number_of_coins, ">=#pending_change": 1, # any amount increase - "pending_coin_removal_count": 1, + "pending_coin_removal_count": number_of_coins, } }, post_block_balance_updates={ 1: { - "confirmed_wallet_balance": -1 * tx_amount, + "confirmed_wallet_balance": -1 * tx_amount * number_of_coins, ">=#spendable_balance": 1, # any amount increase ">=#max_send_amount": 1, # any amount increase "<=#pending_change": -1, # any amount decrease - "pending_coin_removal_count": -1, + "pending_coin_removal_count": -number_of_coins, } }, ), @@ -218,57 +222,19 @@ async def test_wallet_clawback_claim_auto(self, wallet_environments: WalletTestF ] ) - await time_out_assert(20, wsm.coin_store.count_small_unspent, 1, 1000, CoinType.CLAWBACK) - await time_out_assert(20, wsm_1.coin_store.count_small_unspent, 1, 1000, CoinType.CLAWBACK) - - [tx2] = await wallet.generate_signed_transaction( - uint64(tx_amount), - normal_puzhash, - DEFAULT_TX_CONFIG, - uint64(0), - puzzle_decorator_override=[{"decorator": "CLAWBACK", "clawback_timelock": 10}], - ) - [tx2] = await wallet.wallet_state_manager.add_pending_transactions([tx2]) - - await wallet_environments.process_pending_states( - [ - WalletStateTransition( - pre_block_balance_updates={ - 1: { - "unconfirmed_wallet_balance": -1 * tx_amount, - "<=#spendable_balance": -1 * tx_amount, - "<=#max_send_amount": -1 * tx_amount, - ">=#pending_change": 1, # any amount increase - "pending_coin_removal_count": 1, - } - }, - post_block_balance_updates={ - 1: { - "confirmed_wallet_balance": -1 * tx_amount, - ">=#spendable_balance": 1, # any amount increase - ">=#max_send_amount": 1, # any amount increase - "<=#pending_change": -1, # any amount decrease - "pending_coin_removal_count": -1, - } - }, - ), - WalletStateTransition( - pre_block_balance_updates={}, - post_block_balance_updates={}, - ), - ] + await time_out_assert(20, wsm.coin_store.count_small_unspent, number_of_coins, tx_amount * 2, CoinType.CLAWBACK) + await time_out_assert( + 20, wsm_1.coin_store.count_small_unspent, number_of_coins, tx_amount * 2, CoinType.CLAWBACK ) - await time_out_assert(20, wsm.coin_store.count_small_unspent, 2, 1000, CoinType.CLAWBACK) - await time_out_assert(20, wsm_1.coin_store.count_small_unspent, 2, 1000, CoinType.CLAWBACK) - [tx3] = await wallet.generate_signed_transaction( + [tx_bad] = await wallet.generate_signed_transaction( uint64(tx_amount), normal_puzhash, DEFAULT_TX_CONFIG, uint64(0), puzzle_decorator_override=[{"decorator": "CLAWBACK", "clawback_timelock": 10}], ) - [tx3] = await wallet.wallet_state_manager.add_pending_transactions([tx3]) + [tx_bad] = await wallet.wallet_state_manager.add_pending_transactions([tx_bad]) await wallet_environments.process_pending_states( [ @@ -299,13 +265,13 @@ async def test_wallet_clawback_claim_auto(self, wallet_environments: WalletTestF ] ) - # Change 3rd coin to test missing metadata case - clawback_coin_id = tx3.additions[0].name() + # Change one coin to test missing metadata case + clawback_coin_id = tx_bad.additions[0].name() coin_record = await wsm_1.coin_store.get_coin_record(clawback_coin_id) assert coin_record is not None await wsm_1.coin_store.add_coin_record(dataclasses.replace(coin_record, metadata=None)) # Claim merkle coin - env_1.node.set_auto_claim(AutoClaimSettings(enabled=True)) + env_1.node.set_auto_claim(AutoClaimSettings(enabled=True, batch_size=uint16(2))) # Trigger auto claim await wallet_environments.process_pending_states( [ @@ -315,9 +281,10 @@ async def test_wallet_clawback_claim_auto(self, wallet_environments: WalletTestF # After auto claim is set, the next block will trigger submission of clawback claims post_block_balance_updates={ 1: { - "unconfirmed_wallet_balance": 1000, - "pending_change": 1000, # This is a little weird but I think intentional and correct - "pending_coin_removal_count": 2, + "unconfirmed_wallet_balance": tx_amount * number_of_coins, + "pending_change": tx_amount + * number_of_coins, # This is a little weird but I think intentional and correct + "pending_coin_removal_count": number_of_coins, } }, ), @@ -330,19 +297,19 @@ async def test_wallet_clawback_claim_auto(self, wallet_environments: WalletTestF pre_block_balance_updates={}, post_block_balance_updates={ 1: { - "confirmed_wallet_balance": 1000, - "spendable_balance": 1000, - "max_send_amount": 1000, - "unspent_coin_count": 2, - "pending_change": -1000, - "pending_coin_removal_count": -2, + "confirmed_wallet_balance": tx_amount * number_of_coins, + "spendable_balance": tx_amount * number_of_coins, + "max_send_amount": tx_amount * number_of_coins, + "unspent_coin_count": number_of_coins, + "pending_change": -tx_amount * number_of_coins, + "pending_coin_removal_count": -1 * number_of_coins, } }, ), ] ) - await time_out_assert(20, wsm.coin_store.count_small_unspent, 1, 1000, CoinType.CLAWBACK) - await time_out_assert(20, wsm_1.coin_store.count_small_unspent, 1, 1000, CoinType.CLAWBACK) + await time_out_assert(20, wsm.coin_store.count_small_unspent, 1, tx_amount * 2, CoinType.CLAWBACK) + await time_out_assert(20, wsm_1.coin_store.count_small_unspent, 1, tx_amount * 2, CoinType.CLAWBACK) @pytest.mark.parametrize( "wallet_environments", diff --git a/chia/wallet/wallet_node.py b/chia/wallet/wallet_node.py index fa909a5d2ce2..5380d8a27283 100644 --- a/chia/wallet/wallet_node.py +++ b/chia/wallet/wallet_node.py @@ -649,9 +649,6 @@ async def _process_new_subscriptions(self) -> None: peer = item.data[1] assert peer is not None await self.new_peak_wallet(new_peak, peer) - # Check if any coin needs auto spending - if self.config.get("auto_claim", {}).get("enabled", False): - await self.wallet_state_manager.auto_claim_coins() else: self.log.debug("Pulled from queue: UNKNOWN %s", item.item_type) assert False @@ -1161,12 +1158,17 @@ async def new_peak_wallet(self, new_peak: NewPeakWallet, peer: WSChiaConnection) if not await self.new_peak_from_untrusted(new_peak_hb, peer): return - if peer.peer_node_id in self.synced_peers: - await self.wallet_state_manager.blockchain.set_finished_sync_up_to(new_peak.height) # todo why do we call this if there was an exception / the sync is not finished async with self.wallet_state_manager.lock: await self.wallet_state_manager.new_peak(new_peak.height) + # Check if any coin needs auto spending + if self.config.get("auto_claim", {}).get("enabled", False): + await self.wallet_state_manager.auto_claim_coins() + + if peer.peer_node_id in self.synced_peers: + await self.wallet_state_manager.blockchain.set_finished_sync_up_to(new_peak.height) + async def new_peak_from_trusted( self, new_peak_hb: HeaderBlock, latest_timestamp: uint64, peer: WSChiaConnection ) -> None: