[LABS-170] Add automatic reorg testing to WalletTestFramework (#21211)

* [LABS-170] Add automatic reorg testing to `WalletTestFramework`

* Actually exempt reorgs when specified

* Comments by @cursor

* Fix CAT wallet tests

* Reorg exempt offer tests

* Fix DID tests

* Fix NFTWallet tests

* Fix RPC tests

* Fix VC tests

* Exempt pool tests

* Add reorg_exempt to top level conftest

* Fix tests (nd maybe break them?)
This commit is contained in:
Matt Hauff
2026-08-10 16:49:16 -05:00
committed by GitHub
parent 41e1289202
commit d90001e3ee
20 changed files with 339 additions and 523 deletions
+1
View File
@@ -1575,4 +1575,5 @@ async def wallet_environments(
for service, rpc_client, wallet_state in zip(wallet_services, wallet_rpc_clients, wallet_states)
],
tx_config,
request.param.get("reorg_exempt", False),
)
+39 -4
View File
@@ -19,6 +19,7 @@ from chia.rpc.rpc_server import RpcServer
from chia.server.server import ChiaServer
from chia.server.start_service import Service
from chia.simulator.full_node_simulator import FullNodeSimulator
from chia.simulator.simulator_protocol import ReorgProtocol
from chia.wallet.transaction_record import LightTransactionRecord
from chia.wallet.util.transaction_type import CLAWBACK_INCOMING_TRANSACTION_TYPES
from chia.wallet.util.tx_config import DEFAULT_TX_CONFIG, TXConfig
@@ -28,6 +29,7 @@ from chia.wallet.wallet_node_api import WalletNodeAPI
from chia.wallet.wallet_request_types import GetWalletBalance
from chia.wallet.wallet_rpc_api import WalletRpcApi
from chia.wallet.wallet_rpc_client import WalletRpcClient
from chia.wallet.wallet_spend_bundle import WalletSpendBundle
from chia.wallet.wallet_state_manager import WalletStateManager
STANDARD_TX_ENDPOINT_ARGS: dict[str, Any] = TransactionEndpoint(
@@ -309,6 +311,7 @@ class WalletTestFramework:
trusted_full_node: bool
environments: list[WalletEnvironment]
tx_config: TXConfig = DEFAULT_TX_CONFIG
reorg_exempt: bool = False
def cmd_tx_endpoint_args(self, env: WalletEnvironment) -> dict[str, Any]:
return {
@@ -343,7 +346,12 @@ class WalletTestFramework:
yield
async def process_pending_states(
self, state_transitions: list[WalletStateTransition], invalid_transactions: list[bytes32] = []
self,
state_transitions: list[WalletStateTransition],
invalid_transactions: list[bytes32] = [],
post_reorg_balance_differences: list[WalletStateTransition] = [],
bundles_to_repush: list[WalletSpendBundle] = [],
reorg_exempt: bool = False,
) -> None:
"""
This is the main entry point for processing state in wallet tests. It does the following things:
@@ -356,6 +364,8 @@ class WalletTestFramework:
6) Checks that if `reuse_puzhash` was set, no new derivations were created
7) Ensures the wallet is in a synced state before progressing to the rest of the test
"""
if len(post_reorg_balance_differences) == 0:
post_reorg_balance_differences = [WalletStateTransition()] * len(self.environments)
# Take note of the number of puzzle hashes if we're supposed to be reusing
if self.tx_config.reuse_puzhash:
puzzle_hash_indexes: list[dict[uint32, int]] = []
@@ -365,7 +375,13 @@ class WalletTestFramework:
ph_indexes[wallet_id] = await env.wallet_state_manager.puzzle_store.get_used_count(wallet_id)
puzzle_hash_indexes.append(ph_indexes)
balances_pre_block_updates: list[dict[uint32, WalletState]] = []
# Check balances after block (and reorg)
for reorg_status in ("no",) if reorg_exempt or self.reorg_exempt else ("before", "during"):
pending_txs: list[list[LightTransactionRecord]] = []
if reorg_status == "during":
for bundle in bundles_to_repush:
await self.full_node_rpc_client.push_tx(bundle)
peak = self.full_node.full_node.blockchain.get_peak_height()
assert peak is not None
# Check balances prior to block
@@ -383,17 +399,21 @@ class WalletTestFramework:
for i, (env, transition) in enumerate(zip(self.environments, state_transitions)):
try:
async with env.wallet_state_manager.db_wrapper.reader_no_transaction():
if reorg_status == "during":
env.wallet_states = balances_pre_block_updates[i]
await env.change_balances(post_reorg_balance_differences[i].pre_block_balance_updates)
else:
await env.change_balances(transition.pre_block_balance_updates)
balances_pre_block_updates.append(env.wallet_states)
await env.check_balances(transition.pre_block_additional_balance_info)
except Exception:
raise ValueError(f"Error with env index {i}")
raise ValueError(f"Error with env index {i} - {reorg_status} reorg check")
except Exception as e:
raise ValueError(f"Error before block was farmed: {e}") from e
# Farm block
await self.full_node.farm_blocks_to_puzzlehash(count=1, guarantee_transaction_blocks=True)
# Check balances after block
try:
for i, (env, local_pending_txs) in enumerate(zip(self.environments, pending_txs)):
await self.full_node.wait_for_wallet_synced(
@@ -410,9 +430,17 @@ class WalletTestFramework:
try:
async with env.wallet_state_manager.db_wrapper.reader_no_transaction():
await env.change_balances(transition.post_block_balance_updates)
if reorg_status == "during":
await env.change_balances(post_reorg_balance_differences[i].post_block_balance_updates)
await env.check_balances(transition.post_block_additional_balance_info)
if reorg_status == "before":
for id, balance_updates in transition.post_block_balance_updates.items():
if balance_updates.get("init", False):
balances_pre_block_updates[i][env.dealias_wallet_id(id)] = WalletState(
balance=Balance()
)
except Exception:
raise ValueError(f"Error with env {i}")
raise ValueError(f"Error with env {i} - {reorg_status} reorg check")
except Exception as e:
raise ValueError(f"Error after block was farmed: {e}") from e
@@ -428,6 +456,13 @@ class WalletTestFramework:
f"ENV-{i} TXs not confirmed: {[tx.to_json_dict() for tx in unconfirmed if tx in txs]}"
)
if reorg_status == "before":
height = self.full_node.full_node.blockchain.get_peak_height()
assert height is not None
await self.full_node.reorg_from_index_to_new_index(
ReorgProtocol(uint32(height - 1), uint32(height + 1), bytes32.zeros, None)
)
# Finally, check that the number of puzzle hashes did or did not increase by the specified amount
if self.tx_config.reuse_puzhash:
for env, ph_indexes_before in zip(self.environments, puzzle_hash_indexes):
+9
View File
@@ -76,6 +76,7 @@ class StateUrlCase:
{
"num_environments": 1,
"blocks_needed": [1],
"reorg_exempt": True,
}
],
indirect=True,
@@ -229,6 +230,7 @@ async def test_plotnft_cli_create_errors(
{
"num_environments": 1,
"blocks_needed": [1],
"reorg_exempt": True,
}
],
indirect=True,
@@ -309,6 +311,7 @@ async def test_plotnft_cli_show(
{
"num_environments": 1,
"blocks_needed": [1],
"reorg_exempt": True,
}
],
indirect=True,
@@ -381,6 +384,7 @@ async def test_plotnft_cli_show_with_farmer(
{
"num_environments": 1,
"blocks_needed": [10],
"reorg_exempt": True,
}
],
indirect=True,
@@ -472,6 +476,7 @@ async def test_plotnft_cli_leave(
{
"num_environments": 1,
"blocks_needed": [10],
"reorg_exempt": True,
}
],
indirect=True,
@@ -704,6 +709,7 @@ async def test_plotnft_cli_join(
{
"num_environments": 1,
"blocks_needed": [10],
"reorg_exempt": True,
}
],
indirect=True,
@@ -829,6 +835,7 @@ async def test_plotnft_cli_claim(wallet_environments: WalletTestFramework, versi
"num_environments": 1,
"blocks_needed": [10],
"reuse_puzhash": False,
"reorg_exempt": True,
}
],
indirect=True,
@@ -918,6 +925,7 @@ async def test_plotnft_cli_inspect(
{
"num_environments": 1,
"blocks_needed": [10],
"reorg_exempt": True,
}
],
indirect=True,
@@ -990,6 +998,7 @@ async def test_plotnft_cli_change_payout(
{
"num_environments": 1,
"blocks_needed": [10],
"reorg_exempt": True,
}
],
indirect=True,
+3
View File
@@ -1171,6 +1171,7 @@ class TestPoolWalletRpc:
{
"num_environments": 1,
"blocks_needed": [10],
"reorg_exempt": True,
}
],
indirect=True,
@@ -1359,6 +1360,7 @@ class TestPoolWalletRpc:
{
"num_environments": 1,
"blocks_needed": [10],
"reorg_exempt": True,
}
],
indirect=True,
@@ -1401,6 +1403,7 @@ class TestPoolWalletRpc:
"blocks_needed": [10],
"trusted": True,
"reuse_puzhash": False,
"reorg_exempt": True,
}
],
indirect=True,
@@ -1320,7 +1320,25 @@ async def test_cat_hint(wallet_environments: WalletTestFramework, wallet_type: t
}
),
),
]
],
post_reorg_balance_differences=[
WalletStateTransition(),
WalletStateTransition()
if autodiscovery
# after the reorg we'll already have discovered that these CATs belong to us
else WalletStateTransition(
pre_block_balance_updates={
"cat": {
"confirmed_wallet_balance": 60,
"unconfirmed_wallet_balance": 60,
"spendable_balance": 60,
"max_send_amount": 60,
"unspent_coin_count": 1,
}
},
post_block_balance_updates={}, # the "init" means the balances will be overidden
),
],
)
cat_wallet_2 = wallet_node_2.wallet_state_manager.wallets[uint32(2)]
@@ -1522,7 +1540,21 @@ async def test_cat_change_detection(wallet_environments: WalletTestFramework, wa
}
},
)
]
],
bundles_to_repush=[eve_spend],
# We're going to remember we received the CAT through the reorg
post_reorg_balance_differences=[
WalletStateTransition(
pre_block_balance_updates={
"cat": {
"unconfirmed_wallet_balance": 5,
"pending_change": 5,
"pending_coin_removal_count": 1,
}
},
post_block_balance_updates={}, # the "init" means the balances will be overidden
)
],
)
assert not full_node_api.full_node.subscriptions.has_puzzle_subscription(puzzlehash_unhardened)
@@ -1638,7 +1670,8 @@ async def test_cat_melt_balance(wallet_environments: WalletTestFramework) -> Non
},
},
)
]
],
bundles_to_repush=[spend_to_wallet],
)
cat_wallet = env.wallet_state_manager.wallets[uint32(2)]
@@ -1688,7 +1721,8 @@ async def test_cat_melt_balance(wallet_environments: WalletTestFramework) -> Non
},
},
)
]
],
bundles_to_repush=[signed_spend],
)
+65 -8
View File
@@ -63,42 +63,94 @@ async def get_trade_and_status(trade_manager: TradeManager, trade: TradeRecord)
"wallet_environments,credential_restricted,active_softfork_height",
[
(
{"num_environments": 2, "trusted": True, "blocks_needed": [1, 1], "reuse_puzhash": True},
{
"num_environments": 2,
"trusted": True,
"blocks_needed": [1, 1],
"reuse_puzhash": True,
"reorg_exempt": True,
},
True,
SOFTFORK_HEIGHTS[0],
),
(
{"num_environments": 2, "trusted": True, "blocks_needed": [1, 1], "reuse_puzhash": True},
{
"num_environments": 2,
"trusted": True,
"blocks_needed": [1, 1],
"reuse_puzhash": True,
"reorg_exempt": True,
},
False,
SOFTFORK_HEIGHTS[0],
),
(
{"num_environments": 2, "trusted": True, "blocks_needed": [1, 1], "reuse_puzhash": False},
{
"num_environments": 2,
"trusted": True,
"blocks_needed": [1, 1],
"reuse_puzhash": False,
"reorg_exempt": True,
},
True,
SOFTFORK_HEIGHTS[0],
),
(
{"num_environments": 2, "trusted": False, "blocks_needed": [1, 1], "reuse_puzhash": True},
{
"num_environments": 2,
"trusted": False,
"blocks_needed": [1, 1],
"reuse_puzhash": True,
"reorg_exempt": True,
},
True,
SOFTFORK_HEIGHTS[0],
),
(
{"num_environments": 2, "trusted": False, "blocks_needed": [1, 1], "reuse_puzhash": False},
{
"num_environments": 2,
"trusted": False,
"blocks_needed": [1, 1],
"reuse_puzhash": False,
"reorg_exempt": True,
},
False,
SOFTFORK_HEIGHTS[0],
),
(
{"num_environments": 2, "trusted": False, "blocks_needed": [1, 1], "reuse_puzhash": True},
{
"num_environments": 2,
"trusted": False,
"blocks_needed": [1, 1],
"reuse_puzhash": True,
"reorg_exempt": True,
},
False,
SOFTFORK_HEIGHTS[0],
),
(
{"num_environments": 2, "trusted": False, "blocks_needed": [1, 1], "reuse_puzhash": False},
{
"num_environments": 2,
"trusted": False,
"blocks_needed": [1, 1],
"reuse_puzhash": False,
"reorg_exempt": True,
},
True,
SOFTFORK_HEIGHTS[0],
),
*(
({"num_environments": 2, "trusted": True, "blocks_needed": [1, 1], "reuse_puzhash": False}, False, height)
(
{
"num_environments": 2,
"trusted": True,
"blocks_needed": [1, 1],
"reuse_puzhash": False,
"reorg_exempt": True,
},
False,
height,
)
for height in SOFTFORK_HEIGHTS
),
],
@@ -1651,6 +1703,7 @@ async def test_cat_trades(
{
"num_environments": 2,
"blocks_needed": [2, 1],
"reorg_exempt": True,
}
],
indirect=True,
@@ -1963,6 +2016,7 @@ async def test_trade_cancellation(wallet_environments: WalletTestFramework, wall
{
"num_environments": 3,
"blocks_needed": [2, 1, 1],
"reorg_exempt": True,
}
],
indirect=True,
@@ -2160,6 +2214,7 @@ async def test_trade_conflict(wallet_environments: WalletTestFramework, wallet_t
{
"num_environments": 2,
"blocks_needed": [1, 1],
"reorg_exempt": True,
}
],
indirect=True,
@@ -2286,6 +2341,7 @@ async def test_trade_bad_spend(
{
"num_environments": 2,
"blocks_needed": [1, 1],
"reorg_exempt": True,
}
],
indirect=True,
@@ -2433,6 +2489,7 @@ async def test_trade_high_fee(wallet_environments: WalletTestFramework, wallet_t
{
"num_environments": 2,
"blocks_needed": [1, 1],
"reorg_exempt": True,
}
],
indirect=True,
+1
View File
@@ -277,4 +277,5 @@ async def wallet_environments(
for service, rpc_client, wallet_state in zip(wallet_services, wallet_rpc_clients, wallet_states)
],
tx_config,
request.param.get("reorg_exempt", False),
)
@@ -39,7 +39,9 @@ def get_parent_branch(value: bytes32, proof: tuple[int, list[bytes32]]) -> tuple
@pytest.mark.limit_consensus_modes
@pytest.mark.parametrize("wallet_environments", [{"num_environments": 2, "blocks_needed": [2, 2]}], indirect=True)
@pytest.mark.parametrize(
"wallet_environments", [{"num_environments": 2, "blocks_needed": [2, 2], "reorg_exempt": True}], indirect=True
)
@pytest.mark.anyio
async def test_dl_offers(wallet_environments: WalletTestFramework) -> None:
env_maker = wallet_environments.environments[0]
@@ -353,7 +355,9 @@ async def test_dl_offers(wallet_environments: WalletTestFramework) -> None:
@pytest.mark.limit_consensus_modes
@pytest.mark.parametrize("wallet_environments", [{"num_environments": 1, "blocks_needed": [3]}], indirect=True)
@pytest.mark.parametrize(
"wallet_environments", [{"num_environments": 1, "blocks_needed": [3], "reorg_exempt": True}], indirect=True
)
@pytest.mark.anyio
async def test_dl_offer_cancellation(wallet_environments: WalletTestFramework) -> None:
env_maker = wallet_environments.environments[0]
@@ -579,7 +583,9 @@ async def test_dl_offer_cancellation(wallet_environments: WalletTestFramework) -
@pytest.mark.limit_consensus_modes
@pytest.mark.parametrize("wallet_environments", [{"num_environments": 2, "blocks_needed": [3, 3]}], indirect=True)
@pytest.mark.parametrize(
"wallet_environments", [{"num_environments": 2, "blocks_needed": [3, 3], "reorg_exempt": True}], indirect=True
)
@pytest.mark.anyio
async def test_multiple_dl_offers(wallet_environments: WalletTestFramework) -> None:
env_maker = wallet_environments.environments[0]
+22 -2
View File
@@ -469,7 +469,14 @@ async def test_did_find_lost_did(wallet_environments: WalletTestFramework, capsy
},
},
),
]
],
# Due to an ephemeral spend, sync will pick up an extra TX here and thus bump the pending coin removal count
post_reorg_balance_differences=[
WalletStateTransition(
pre_block_balance_updates={"did_found": {"pending_coin_removal_count": 1}},
post_block_balance_updates={"did_found": {"pending_coin_removal_count": -1}},
)
],
)
coin = await did_wallet.get_coin()
@@ -587,7 +594,10 @@ async def test_did_transfer(wallet_environments: WalletTestFramework, capsys: py
"did": {"init": True, "confirmed_wallet_balance": 101, "set_remainder": True},
},
),
]
],
# TODO: there's a bug here where the user store autoincrement means this has a new ID after deleted
# Instead of 2, it becomes 3 because there was a 2 at some point (is my best guess)
reorg_exempt=True,
)
did_wallets = list(
@@ -1046,7 +1056,17 @@ async def test_update_metadata(wallet_environments: WalletTestFramework, capsys:
"did": {"confirmed_wallet_balance": 0, "set_remainder": True},
},
),
],
post_reorg_balance_differences=[
WalletStateTransition(
pre_block_balance_updates={"did": {"pending_coin_removal_count": 1}},
post_block_balance_updates={}, # set_remainder takes care of it
),
WalletStateTransition(),
]
# TODO: figure out why the reuse_puzhash being off is the only reason this happens
if not wallet_environments.tx_config.reuse_puzhash
else [],
)
assert get_parent_num(did_wallet_1) == parent_num + 2
@@ -29,7 +29,9 @@ async def get_nft_count(wallet: NFTWallet) -> int:
@pytest.mark.limit_consensus_modes
@pytest.mark.parametrize("wallet_environments", [{"num_environments": 2, "blocks_needed": [1, 1]}], indirect=True)
@pytest.mark.parametrize(
"wallet_environments", [{"num_environments": 2, "blocks_needed": [1, 1], "reorg_exempt": True}], indirect=True
)
@pytest.mark.parametrize("zero_royalties", [True, False])
@pytest.mark.anyio
async def test_nft_offer_sell_nft(wallet_environments: WalletTestFramework, zero_royalties: bool) -> None:
@@ -255,7 +257,9 @@ async def test_nft_offer_sell_nft(wallet_environments: WalletTestFramework, zero
@pytest.mark.limit_consensus_modes
@pytest.mark.parametrize("wallet_environments", [{"num_environments": 2, "blocks_needed": [1, 1]}], indirect=True)
@pytest.mark.parametrize(
"wallet_environments", [{"num_environments": 2, "blocks_needed": [1, 1], "reorg_exempt": True}], indirect=True
)
@pytest.mark.parametrize("zero_royalties", [True, False])
@pytest.mark.anyio
async def test_nft_offer_request_nft(wallet_environments: WalletTestFramework, zero_royalties: bool) -> None:
@@ -485,7 +489,9 @@ async def test_nft_offer_request_nft(wallet_environments: WalletTestFramework, z
@pytest.mark.limit_consensus_modes
@pytest.mark.parametrize("wallet_environments", [{"num_environments": 2, "blocks_needed": [1, 1]}], indirect=True)
@pytest.mark.parametrize(
"wallet_environments", [{"num_environments": 2, "blocks_needed": [1, 1], "reorg_exempt": True}], indirect=True
)
@pytest.mark.parametrize("zero_royalties", [True, False])
@pytest.mark.anyio
async def test_nft_offer_sell_did_to_did(wallet_environments: WalletTestFramework, zero_royalties: bool) -> None:
@@ -777,7 +783,9 @@ async def test_nft_offer_sell_did_to_did(wallet_environments: WalletTestFramewor
@pytest.mark.limit_consensus_modes
@pytest.mark.parametrize("wallet_environments", [{"num_environments": 2, "blocks_needed": [1, 1]}], indirect=True)
@pytest.mark.parametrize(
"wallet_environments", [{"num_environments": 2, "blocks_needed": [1, 1], "reorg_exempt": True}], indirect=True
)
@pytest.mark.parametrize("zero_royalties", [True, False])
@pytest.mark.parametrize("wallet_type", [CATWallet, RCATWallet])
@pytest.mark.anyio
@@ -1101,7 +1109,9 @@ async def test_nft_offer_sell_nft_for_cat(
@pytest.mark.limit_consensus_modes
@pytest.mark.parametrize("wallet_environments", [{"num_environments": 2, "blocks_needed": [1, 1]}], indirect=True)
@pytest.mark.parametrize(
"wallet_environments", [{"num_environments": 2, "blocks_needed": [1, 1], "reorg_exempt": True}], indirect=True
)
@pytest.mark.parametrize("test_change", [True, False])
@pytest.mark.parametrize("wallet_type", [CATWallet, RCATWallet])
@pytest.mark.anyio
@@ -1460,7 +1470,9 @@ async def test_nft_offer_request_nft_for_cat(
@pytest.mark.limit_consensus_modes
@pytest.mark.parametrize("wallet_environments", [{"num_environments": 1, "blocks_needed": [2]}], indirect=True)
@pytest.mark.parametrize(
"wallet_environments", [{"num_environments": 1, "blocks_needed": [2], "reorg_exempt": True}], indirect=True
)
@pytest.mark.anyio
async def test_nft_offer_sell_cancel(wallet_environments: WalletTestFramework) -> None:
env_maker = wallet_environments.environments[0]
@@ -1652,7 +1664,14 @@ async def test_nft_offer_sell_cancel(wallet_environments: WalletTestFramework) -
)
@pytest.mark.parametrize(
"wallet_environments",
[{"num_environments": 2, "blocks_needed": [3, 3], "config_overrides": {"automatically_add_unknown_cats": True}}],
[
{
"num_environments": 2,
"blocks_needed": [3, 3],
"config_overrides": {"automatically_add_unknown_cats": True},
"reorg_exempt": True,
}
],
indirect=True,
)
@pytest.mark.parametrize("wallet_type", [CATWallet, RCATWallet])
@@ -26,7 +26,9 @@ async def get_trade_and_status(trade_manager, trade) -> TradeStatus: # type: ig
@pytest.mark.limit_consensus_modes(reason="irrelevant")
@pytest.mark.parametrize("wallet_environments", [{"num_environments": 2, "blocks_needed": [1, 1]}], indirect=True)
@pytest.mark.parametrize(
"wallet_environments", [{"num_environments": 2, "blocks_needed": [1, 1], "reorg_exempt": True}], indirect=True
)
@pytest.mark.anyio
async def test_nft_offer_with_fee(wallet_environments: WalletTestFramework) -> None:
env_0 = wallet_environments.environments[0]
@@ -308,7 +310,9 @@ async def test_nft_offer_with_fee(wallet_environments: WalletTestFramework) -> N
@pytest.mark.limit_consensus_modes(reason="irrelevant")
@pytest.mark.parametrize("wallet_environments", [{"num_environments": 1, "blocks_needed": [1]}], indirect=True)
@pytest.mark.parametrize(
"wallet_environments", [{"num_environments": 1, "blocks_needed": [1], "reorg_exempt": True}], indirect=True
)
@pytest.mark.anyio
async def test_nft_offer_cancellations(wallet_environments: WalletTestFramework) -> None:
env_0 = wallet_environments.environments[0]
@@ -461,7 +465,9 @@ async def test_nft_offer_cancellations(wallet_environments: WalletTestFramework)
@pytest.mark.limit_consensus_modes(reason="irrelevant")
@pytest.mark.parametrize("wallet_environments", [{"num_environments": 2, "blocks_needed": [1, 1]}], indirect=True)
@pytest.mark.parametrize(
"wallet_environments", [{"num_environments": 2, "blocks_needed": [1, 1], "reorg_exempt": True}], indirect=True
)
@pytest.mark.anyio
async def test_nft_offer_with_metadata_update(wallet_environments: WalletTestFramework) -> None:
env_0 = wallet_environments.environments[0]
@@ -696,7 +702,9 @@ async def test_nft_offer_with_metadata_update(wallet_environments: WalletTestFra
@pytest.mark.limit_consensus_modes(reason="irrelevant")
@pytest.mark.parametrize("wallet_environments", [{"num_environments": 2, "blocks_needed": [1, 1]}], indirect=True)
@pytest.mark.parametrize(
"wallet_environments", [{"num_environments": 2, "blocks_needed": [1, 1], "reorg_exempt": True}], indirect=True
)
@pytest.mark.parametrize("wallet_type", [CATWallet, RCATWallet])
@pytest.mark.anyio
async def test_nft_offer_nft_for_cat(wallet_environments: WalletTestFramework, wallet_type: type[CATWallet]) -> None:
@@ -1093,7 +1101,9 @@ async def test_nft_offer_nft_for_cat(wallet_environments: WalletTestFramework, w
@pytest.mark.limit_consensus_modes(reason="irrelevant")
@pytest.mark.parametrize("wallet_environments", [{"num_environments": 2, "blocks_needed": [1, 1]}], indirect=True)
@pytest.mark.parametrize(
"wallet_environments", [{"num_environments": 2, "blocks_needed": [1, 1], "reorg_exempt": True}], indirect=True
)
@pytest.mark.anyio
async def test_nft_offer_nft_for_nft(wallet_environments: WalletTestFramework) -> None:
env_0 = wallet_environments.environments[0]
@@ -1333,7 +1343,9 @@ async def test_nft_offer_nft_for_nft(wallet_environments: WalletTestFramework) -
@pytest.mark.limit_consensus_modes(reason="irrelevant")
@pytest.mark.parametrize("wallet_environments", [{"num_environments": 2, "blocks_needed": [1, 1]}], indirect=True)
@pytest.mark.parametrize(
"wallet_environments", [{"num_environments": 2, "blocks_needed": [1, 1], "reorg_exempt": True}], indirect=True
)
@pytest.mark.parametrize("wallet_type", [CATWallet, RCATWallet])
@pytest.mark.anyio
async def test_nft_offer_nft0_and_xch_for_cat(
@@ -1595,7 +1595,10 @@ async def test_nft_transfer_nft_with_did(wallet_environments: WalletTestFramewor
}
},
),
]
],
# TODO: there's a bug here where the user store autoincrement means this has a new ID after deleted
# Instead of 2, it becomes 3 because there was a 2 at some point (is my best guess)
reorg_exempt=True,
)
# Transfer NFT, wallet will be deleted
@@ -90,43 +90,6 @@ async def test_plotnft_lifecycle(wallet_environments: WalletTestFramework, self_
uint32(env.wallet_aliases["plotnft"]), required_type=PlotNFT2Wallet
)
# Reorg (creation)
height = wallet_environments.full_node.full_node.blockchain.get_peak_height()
assert height is not None
await wallet_environments.full_node.reorg_from_index_to_new_index(
ReorgProtocol(uint32(height - 1), uint32(height + 1), bytes32.zeros, None)
)
await wallet_environments.full_node.wait_for_wallet_synced(env.node)
await wallet_environments.process_pending_states(
[
WalletStateTransition(
pre_block_balance_updates={
"xch": {
"confirmed_wallet_balance": creation_fee + 1,
"<=#spendable_balance": creation_fee + 1,
"<=#max_send_amount": creation_fee + 1,
">=#pending_change": 0,
">=#pending_coin_removal_count": 1,
">=#unspent_coin_count": 0,
},
"plotnft": {"unspent_coin_count": -1},
},
post_block_balance_updates={
"xch": {
"confirmed_wallet_balance": -(creation_fee + 1),
">=#spendable_balance": 1,
">=#max_send_amount": 1,
"<=#pending_change": 0,
"<=#pending_coin_removal_count": -1,
"<=#unspent_coin_count": 0,
},
"plotnft": {"unspent_coin_count": 1},
},
)
]
)
# (check an error)
with pytest.raises(ValueError, match=re.escape("`leave_pool` called on a non-pooling or exiting PlotNFT")):
async with env.wallet_state_manager.new_action_scope(wallet_environments.tx_config, push=True) as action_scope:
@@ -244,49 +207,6 @@ async def test_plotnft_lifecycle(wallet_environments: WalletTestFramework, self_
):
await plotnft_wallet.claim_rewards(action_scope=action_scope)
# Reorg (claim rewards)
height = wallet_environments.full_node.full_node.blockchain.get_peak_height()
assert height is not None
await wallet_environments.full_node.reorg_from_index_to_new_index(
ReorgProtocol(uint32(height - 1), uint32(height + 1), bytes32.zeros, None)
)
await wallet_environments.full_node.wait_for_wallet_synced(env.node)
await wallet_environments.process_pending_states(
[
WalletStateTransition(
pre_block_balance_updates={
"xch": {
"confirmed_wallet_balance": -amount_to_succeed_in_claiming,
"unconfirmed_wallet_balance": -amount_to_succeed_in_claiming,
"spendable_balance": -amount_to_succeed_in_claiming,
"max_send_amount": -amount_to_succeed_in_claiming,
"unspent_coin_count": -1,
},
"plotnft": {
"confirmed_wallet_balance": REWARDS_GAINED,
"pending_coin_removal_count": NUM_REWARDS_FARMED + 1,
"unspent_coin_count": NUM_REWARDS_FARMED,
},
},
post_block_balance_updates={
"xch": {
"confirmed_wallet_balance": amount_to_succeed_in_claiming,
"unconfirmed_wallet_balance": amount_to_succeed_in_claiming,
"spendable_balance": amount_to_succeed_in_claiming,
"max_send_amount": amount_to_succeed_in_claiming,
"unspent_coin_count": 1,
},
"plotnft": {
"confirmed_wallet_balance": -REWARDS_GAINED,
"pending_coin_removal_count": -NUM_REWARDS_FARMED - 1,
"unspent_coin_count": -NUM_REWARDS_FARMED,
},
},
)
]
)
# JOIN POOL
joining_fee = uint64(1_000)
@@ -332,43 +252,6 @@ async def test_plotnft_lifecycle(wallet_environments: WalletTestFramework, self_
== "https://daurl.com"
)
# Reorg (join pool)
height = wallet_environments.full_node.full_node.blockchain.get_peak_height()
assert height is not None
await wallet_environments.full_node.reorg_from_index_to_new_index(
ReorgProtocol(uint32(height - 1), uint32(height + 1), bytes32.zeros, None)
)
await wallet_environments.full_node.wait_for_wallet_synced(env.node)
await wallet_environments.process_pending_states(
[
WalletStateTransition(
pre_block_balance_updates={
"xch": {
"confirmed_wallet_balance": joining_fee,
"<=#spendable_balance": -1,
"<=#max_send_amount": -1,
">=#pending_change": 0,
">=#pending_coin_removal_count": 1,
">=#unspent_coin_count": 0,
},
"plotnft": {"pending_coin_removal_count": 1},
},
post_block_balance_updates={
"xch": {
"confirmed_wallet_balance": -joining_fee,
">=#spendable_balance": 1,
">=#max_send_amount": 1,
"<=#pending_change": 0,
"<=#pending_coin_removal_count": -1,
"<=#unspent_coin_count": 0,
},
"plotnft": {"pending_coin_removal_count": -1},
},
)
]
)
# RECEIVE REWARDS (while pooling)
EXTRA_POOLING_REWARDS = 2
await wallet_environments.full_node.farm_blocks_to_puzzlehash(
@@ -429,7 +312,15 @@ async def test_plotnft_lifecycle(wallet_environments: WalletTestFramework, self_
}
},
)
]
],
# when we reorg, we'll still remember that we saw an attempt to spend our plotnft
post_reorg_balance_differences=[
WalletStateTransition(
pre_block_balance_updates={"plotnft": {"pending_coin_removal_count": 2}},
post_block_balance_updates={"plotnft": {"pending_coin_removal_count": -2}},
)
],
bundles_to_repush=[WalletSpendBundle(coin_spends, G2Element())],
)
# LEAVE POOL (to another)
@@ -469,33 +360,6 @@ async def test_plotnft_lifecycle(wallet_environments: WalletTestFramework, self_
async with env.wallet_state_manager.new_action_scope(wallet_environments.tx_config, push=True) as action_scope:
await plotnft_wallet.leave_pool(action_scope=action_scope)
# Reorg (leave pool to another)
height = wallet_environments.full_node.full_node.blockchain.get_peak_height()
assert height is not None
await wallet_environments.full_node.reorg_from_index_to_new_index(
ReorgProtocol(uint32(height - 1), uint32(height + 1), bytes32.zeros, None)
)
await wallet_environments.full_node.wait_for_wallet_synced(env.node)
await wallet_environments.process_pending_states(
[
WalletStateTransition(
pre_block_balance_updates={
"xch": {},
"plotnft": {
"pending_coin_removal_count": 1,
},
},
post_block_balance_updates={
"xch": {},
"plotnft": {
"pending_coin_removal_count": -1,
},
},
)
]
)
# FINISH LEAVING (to new pool)
plotnft = await plotnft_wallet.get_current_plotnft()
await wallet_environments.full_node.farm_blocks_to_puzzlehash(
@@ -581,47 +445,6 @@ async def test_plotnft_lifecycle(wallet_environments: WalletTestFramework, self_
):
await plotnft_wallet.claim_rewards(action_scope=action_scope)
# Reorg (leave pool)
height = wallet_environments.full_node.full_node.blockchain.get_peak_height()
assert height is not None
await wallet_environments.full_node.reorg_from_index_to_new_index(
ReorgProtocol(uint32(height - 1), uint32(height + 1), bytes32.zeros, None)
)
await wallet_environments.full_node.wait_for_wallet_synced(env.node)
await wallet_environments.process_pending_states(
[
WalletStateTransition(
pre_block_balance_updates={
"xch": {
"confirmed_wallet_balance": leave_fee,
"<=#spendable_balance": -1,
"<=#max_send_amount": -1,
">=#pending_change": 0,
">=#pending_coin_removal_count": 1,
">=#unspent_coin_count": 0,
},
"plotnft": {
"pending_coin_removal_count": 1,
},
},
post_block_balance_updates={
"xch": {
"confirmed_wallet_balance": -leave_fee,
">=#spendable_balance": 1,
">=#max_send_amount": 1,
"<=#pending_change": 0,
"<=#pending_coin_removal_count": -1,
"<=#unspent_coin_count": 0,
},
"plotnft": {
"pending_coin_removal_count": -1,
},
},
)
]
)
# LOSE REWARDS (while leaving)
plotnft = await plotnft_wallet.get_current_plotnft()
[pool_reward] = await env.wallet_state_manager.plotnft2_store.get_pool_rewards(plotnft_id=plotnft_wallet.plotnft_id)
@@ -642,7 +465,8 @@ async def test_plotnft_lifecycle(wallet_environments: WalletTestFramework, self_
}
},
)
]
],
bundles_to_repush=[WalletSpendBundle(coin_spends, G2Element())],
)
# FINISH LEAVING
@@ -684,47 +508,6 @@ async def test_plotnft_lifecycle(wallet_environments: WalletTestFramework, self_
]
)
# Reorg (finish leaving)
height = wallet_environments.full_node.full_node.blockchain.get_peak_height()
assert height is not None
await wallet_environments.full_node.reorg_from_index_to_new_index(
ReorgProtocol(uint32(height - 1), uint32(height + 1), bytes32.zeros, None)
)
await wallet_environments.full_node.wait_for_wallet_synced(env.node)
await wallet_environments.process_pending_states(
[
WalletStateTransition(
pre_block_balance_updates={
"xch": {
"confirmed_wallet_balance": finish_leaving_fee,
"<=#spendable_balance": 0,
"<=#max_send_amount": 0,
">=#pending_change": 0,
">=#pending_coin_removal_count": 1,
">=#unspent_coin_count": 0,
},
"plotnft": {
"pending_coin_removal_count": 1,
},
},
post_block_balance_updates={
"xch": {
"confirmed_wallet_balance": -finish_leaving_fee,
">=#spendable_balance": 0,
">=#max_send_amount": 0,
"<=#pending_change": 0,
"<=#pending_coin_removal_count": -1,
"<=#unspent_coin_count": 0,
},
"plotnft": {
"pending_coin_removal_count": -1,
},
},
)
]
)
# Resync start
env.node._close()
await env.node._await_closed()
@@ -762,6 +545,7 @@ async def test_plotnft_lifecycle(wallet_environments: WalletTestFramework, self_
{
"num_environments": 1,
"blocks_needed": [1],
"reorg_exempt": True,
}
],
indirect=True,
+10 -2
View File
@@ -1563,6 +1563,7 @@ async def test_cat_endpoints(wallet_environments: WalletTestFramework, wallet_ty
{
"num_environments": 2,
"blocks_needed": [1, 1],
"reorg_exempt": True,
}
],
indirect=True,
@@ -2299,7 +2300,11 @@ async def test_did_endpoints(wallet_environments: WalletTestFramework, capsys: p
},
),
WalletStateTransition(),
]
],
post_reorg_balance_differences=[
WalletStateTransition({"did": {"set_remainder": True}}),
WalletStateTransition(),
],
)
# Transfer DID
@@ -2328,7 +2333,10 @@ async def test_did_endpoints(wallet_environments: WalletTestFramework, capsys: p
"did": {"init": True, "set_remainder": True},
}
),
]
],
# TODO: there's a bug here where the user store autoincrement means this has a new ID after deleted
# Instead of 2, it becomes 3 because there was a 2 at some point (is my best guess)
reorg_exempt=True,
)
async def num_wallets() -> int:
+4 -6
View File
@@ -263,14 +263,11 @@ async def test_p2dohp_wallet_signer_protocol(wallet_environments: WalletTestFram
)
).signed_transactions
await wallet_rpc.submit_transactions(SubmitTransactions(signed_transactions=signed_txs))
await wallet_environments.full_node.wait_bundle_ids_in_mempool(
[
WalletSpendBundle(
bundle = WalletSpendBundle(
[spend.as_coin_spend() for tx in signed_txs for spend in tx.transaction_info.spends],
G2Element.from_bytes(signing_responses[0].signature),
).name()
]
)
await wallet_environments.full_node.wait_bundle_ids_in_mempool([bundle.name()])
await wallet_environments.process_pending_states(
[
@@ -287,7 +284,8 @@ async def test_p2dohp_wallet_signer_protocol(wallet_environments: WalletTestFram
},
},
),
]
],
bundles_to_repush=[bundle],
)
# And test that we can get compressed versions if we want
+2 -189
View File
@@ -427,7 +427,8 @@ class TestWalletSimulator:
}
},
),
]
],
reorg_exempt=True,
)
await wallet_environments.process_pending_states(
[
@@ -890,194 +891,6 @@ class TestWalletSimulator:
assert txs_response.transactions[0].confirmed
assert txs_response.transactions[1].confirmed
@pytest.mark.parametrize(
"wallet_environments",
[{"num_environments": 2, "blocks_needed": [1, 1], "reuse_puzhash": True}],
indirect=True,
)
@pytest.mark.limit_consensus_modes(reason="irrelevant")
@pytest.mark.anyio
async def test_wallet_clawback_reorg(self, wallet_environments: WalletTestFramework) -> None:
full_node_api = wallet_environments.full_node
env = wallet_environments.environments[0]
env_2 = wallet_environments.environments[1]
wsm = env.wallet_state_manager
wsm_2 = env_2.wallet_state_manager
tx_amount = 500
async with wsm_2.new_action_scope(DEFAULT_TX_CONFIG, push=True) as action_scope:
normal_puzhash = await action_scope.get_puzzle_hash(wsm_2)
# Transfer to normal wallet
await env.rpc_client.send_transaction(
SendTransaction(
wallet_id=env.xch_wallet.id(),
amount=uint64(tx_amount),
address=env.wallet_state_manager.encode_puzzle_hash(normal_puzhash),
puzzle_decorator=[ClawbackPuzzleDecoratorOverride(decorator="CLAWBACK", clawback_timelock=uint64(5))],
push=True,
),
wallet_environments.tx_config,
)
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={},
),
]
)
# Check merkle coins
await time_out_assert(20, wsm.coin_store.count_small_unspent, 1, 1000, CoinType.CLAWBACK)
await time_out_assert(20, wsm_2.coin_store.count_small_unspent, 1, 1000, CoinType.CLAWBACK)
# Reorg before claim
# Test Reorg mint
height = full_node_api.full_node.blockchain.get_peak_height()
assert height is not None
await full_node_api.reorg_from_index_to_new_index(
ReorgProtocol(uint32(height - 2), uint32(height + 1), bytes32.zeros, None)
)
await time_out_assert(20, wsm.coin_store.count_small_unspent, 0, 1000, CoinType.CLAWBACK)
await time_out_assert(20, wsm_2.coin_store.count_small_unspent, 0, 1000, CoinType.CLAWBACK)
await wallet_environments.process_pending_states(
[
WalletStateTransition(
pre_block_balance_updates={
1: {
"confirmed_wallet_balance": tx_amount, # confirmed balance comes back
# clawback transaction is now outstanding
"<=#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, 1, 1000, CoinType.CLAWBACK)
await time_out_assert(20, wsm_2.coin_store.count_small_unspent, 1, 1000, CoinType.CLAWBACK)
# Claim merkle coin
await env_2.rpc_client.set_auto_claim(AutoClaimSettings(enabled=True))
# clawback merkle coin
await wallet_environments.process_pending_states(
[
WalletStateTransition(),
WalletStateTransition(
pre_block_balance_updates={},
# After auto claim is set, the next block will trigger submission of clawback claims
post_block_balance_updates={
1: {
"unconfirmed_wallet_balance": tx_amount,
"pending_change": tx_amount, # This is a little weird but I think intentional and correct
"pending_coin_removal_count": 1,
}
},
),
]
)
await wallet_environments.process_pending_states(
[
WalletStateTransition(),
WalletStateTransition(
pre_block_balance_updates={},
post_block_balance_updates={
1: {
"confirmed_wallet_balance": tx_amount,
"spendable_balance": tx_amount,
"max_send_amount": tx_amount,
"unspent_coin_count": 1,
"pending_change": -1 * tx_amount,
"pending_coin_removal_count": -1,
}
},
),
]
)
await time_out_assert(20, wsm.coin_store.count_small_unspent, 0, 1000, CoinType.CLAWBACK)
await time_out_assert(20, wsm_2.coin_store.count_small_unspent, 0, 1000, CoinType.CLAWBACK)
# Reorg after claim
height = full_node_api.full_node.blockchain.get_peak_height()
assert height is not None
await full_node_api.reorg_from_index_to_new_index(
ReorgProtocol(uint32(height - 1), uint32(height + 1), bytes32.zeros, None)
)
await time_out_assert(20, wsm.coin_store.count_small_unspent, 1, 1000, CoinType.CLAWBACK)
await time_out_assert(20, wsm_2.coin_store.count_small_unspent, 1, 1000, CoinType.CLAWBACK)
await wallet_environments.process_pending_states(
[
WalletStateTransition(
pre_block_balance_updates={},
post_block_balance_updates={},
),
WalletStateTransition(
pre_block_balance_updates={
1: {
"confirmed_wallet_balance": -1 * tx_amount,
"spendable_balance": -1 * tx_amount,
"max_send_amount": -1 * tx_amount,
"unspent_coin_count": -1,
"pending_change": tx_amount,
"pending_coin_removal_count": 1,
}
},
post_block_balance_updates={
1: {
"confirmed_wallet_balance": tx_amount,
"spendable_balance": tx_amount,
"max_send_amount": tx_amount,
"unspent_coin_count": 1,
"pending_change": -1 * tx_amount,
"pending_coin_removal_count": -1,
}
},
),
]
)
await time_out_assert(20, wsm.coin_store.count_small_unspent, 0, 1000, CoinType.CLAWBACK)
await time_out_assert(20, wsm_2.coin_store.count_small_unspent, 0, 1000, CoinType.CLAWBACK)
@pytest.mark.parametrize(
"wallet_environments",
[{"num_environments": 1, "blocks_needed": [1], "trusted": True, "reuse_puzhash": True}],
+5 -2
View File
@@ -295,6 +295,7 @@ class DIDWallet:
async def get_pending_change_balance(self) -> uint64:
unconfirmed_tx = await self.wallet_state_manager.tx_store.get_unconfirmed_for_wallet(self.id())
addition_amount = 0
counted_additions = set()
for record in unconfirmed_tx:
our_spend = False
@@ -314,9 +315,10 @@ class DIDWallet:
if len(memos) > 0 and len(memos[0]) == 32
}
if (await self.wallet_state_manager.does_coin_belong_to_wallet(coin, self.id(), hint_dict)) and (
coin not in record.removals
coin not in record.removals and coin not in counted_additions
):
addition_amount += coin.amount
counted_additions.add(coin)
return uint64(addition_amount)
@@ -1034,7 +1036,8 @@ class DIDWallet:
async def add_parent(self, name: bytes32, parent: LineageProof | None) -> None:
self.log.info(f"Adding parent {name}: {parent}")
current_list = self.did_info.parent_info.copy()
# coping for not being a dict - thanks streamable!
current_list = [(n, p) for n, p in self.did_info.parent_info if n != name]
current_list.append((name, parent))
did_info = DIDInfo(
origin_coin=self.did_info.origin_coin,
+1 -1
View File
@@ -791,7 +791,7 @@ class CRCATWallet(CATWallet):
wallet_id=self.id(),
sent_to=[],
trade_id=None,
type=uint32(TransactionType.INCOMING_TX.value),
type=uint32(TransactionType.OUTGOING_TX.value),
name=claim_bundle.name(),
memos=compute_memos(claim_bundle),
valid_times=parse_timelock_info(extra_conditions),
+6 -1
View File
@@ -115,6 +115,7 @@ class Wallet:
unconfirmed_tx: list[TransactionRecord] = await self.wallet_state_manager.tx_store.get_unconfirmed_for_wallet(
self.id()
)
counted_additions = set()
addition_amount = 0
for record in unconfirmed_tx:
@@ -139,8 +140,12 @@ class Wallet:
continue
for coin in record.additions:
if await self.wallet_state_manager.does_coin_belong_to_wallet(coin, self.id()):
if (
await self.wallet_state_manager.does_coin_belong_to_wallet(coin, self.id())
and coin not in counted_additions
):
addition_amount += coin.amount
counted_additions.add(coin)
return uint64(addition_amount)
+5
View File
@@ -881,6 +881,11 @@ class WalletStateManager:
if await self.does_coin_belong_to_wallet(addition, wallet_id, record.hint_dict()):
all_unspent_coins.add(addition)
for record in unconfirmed_tx:
if record.type in CLAWBACK_INCOMING_TRANSACTION_TYPES:
# We do not wish to consider clawback-able funds as unconfirmed.
# That is reserved for when the action to actually claw a tx back or forward is initiated.
continue
for removal in record.removals:
if (
await self.does_coin_belong_to_wallet(removal, wallet_id, record.hint_dict())