Skip to content

Commit

Permalink
[CHIA-738] Add a better clawback auto claim test and fix related issue (
Browse files Browse the repository at this point in the history 
#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
  • Loading branch information
@Quexington
Quexington committed Jul 12, 2024
1 parent 3b4d71e commit 064d34f
Show file tree
Hide file tree
Showing 2 changed files with 48 additions and 79 deletions.
115 changes: 41 additions & 74 deletions chia/_tests/wallet/test_wallet.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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,
}
},
),
Expand All @@ -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(
[
Expand Down Expand Up @@ -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(
[
Expand All @@ -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,
}
},
),
Expand All @@ -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",
Expand Down
12 changes: 7 additions & 5 deletions chia/wallet/wallet_node.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down

0 comments on commit 064d34f

Please sign in to comment.