Add a way to import a SpendSim and SimClient as a contextmanager (#14434)

This commit is contained in:
Matt Hauff 2023-01-31 18:22:06 -07:00 committed by GitHub
parent d741161883
commit 97230ca3b0
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
9 changed files with 106 additions and 162 deletions

View File

@ -1,9 +1,10 @@
from __future__ import annotations
import random
from contextlib import asynccontextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Tuple, Type, TypeVar
from typing import Any, AsyncIterator, Callable, Dict, List, Optional, Tuple, Type, TypeVar
from chia.consensus.block_rewards import calculate_base_farmer_reward, calculate_pool_reward
from chia.consensus.coinbase import create_farmer_coin, create_pool_coin
@ -39,6 +40,20 @@ and is designed so that you could test with it and then swap in a real rpc clien
"""
@asynccontextmanager
async def sim_and_client(
db_path: Optional[Path] = None, defaults: ConsensusConstants = DEFAULT_CONSTANTS, pass_prefarm: bool = True
) -> AsyncIterator[Tuple[SpendSim, SimClient]]:
sim: SpendSim = await SpendSim.create(db_path, defaults)
try:
client: SimClient = SimClient(sim)
if pass_prefarm:
await sim.farm_block()
yield sim, client
finally:
await sim.close()
@streamable
@dataclass(frozen=True)
class SimFullBlock(Streamable):

View File

@ -5,7 +5,7 @@ from typing import List, Optional, Tuple
import pytest
from blspy import AugSchemeMPL, G1Element, G2Element, PrivateKey
from chia.clvm.spend_sim import SimClient, SpendSim
from chia.clvm.spend_sim import SimClient, SpendSim, sim_and_client
from chia.consensus.default_constants import DEFAULT_CONSTANTS
from chia.types.blockchain_format.coin import Coin
from chia.types.blockchain_format.program import Program
@ -80,7 +80,7 @@ class TestSingleton:
@pytest.mark.asyncio
@pytest.mark.parametrize("version", [0, 1])
async def test_singleton_top_layer(self, version):
try:
async with sim_and_client() as (sim, sim_client):
# START TESTS
# Generate starting info
key_lookup = KeyTool()
@ -99,8 +99,6 @@ class TestSingleton:
# Get our starting standard coin created
START_AMOUNT: uint64 = 1023
sim = await SpendSim.create()
sim_client = SimClient(sim)
await sim.farm_block(starting_puzzle.get_tree_hash())
starting_coin: Coin = await sim_client.get_coin_records_by_puzzle_hash(starting_puzzle.get_tree_hash())
starting_coin = starting_coin[0].coin
@ -519,5 +517,3 @@ class TestSingleton:
melted_coin: Coin = (await sim.all_non_reward_coins())[0]
assert melted_coin.puzzle_hash == adapted_puzzle_hash
finally:
await sim.close()

View File

@ -3,7 +3,7 @@ from __future__ import annotations
import pytest
from blspy import G2Element
from chia.clvm.spend_sim import SimClient, SpendSim
from chia.clvm.spend_sim import sim_and_client
from chia.types.blockchain_format.program import Program
from chia.types.blockchain_format.sized_bytes import bytes32
from chia.types.coin_spend import CoinSpend
@ -13,21 +13,17 @@ from chia.types.spend_bundle import SpendBundle
class TestSpendSim:
@pytest.mark.asyncio
async def test_farming(self):
try:
sim = await SpendSim.create()
async with sim_and_client(pass_prefarm=False) as (sim, _):
for i in range(0, 5):
await sim.farm_block()
assert len(sim.blocks) == 5
assert sim.blocks[-1].height == 4
assert sim.block_records[0].reward_claims_incorporated[0].amount == 18375000000000000000
finally:
await sim.close()
@pytest.mark.asyncio
async def test_rewind(self):
try:
sim = await SpendSim.create()
async with sim_and_client() as (sim, _):
for i in range(0, 5):
await sim.farm_block()
@ -35,18 +31,14 @@ class TestSpendSim:
await sim.farm_block()
await sim.rewind(save_height)
assert len(sim.blocks) == 5
assert sim.blocks[-1].height == 4
finally:
await sim.close()
assert len(sim.blocks) == 6
assert sim.blocks[-1].height == 5
class TestSimClient:
@pytest.mark.asyncio
async def test_all_endpoints(self):
try:
sim = await SpendSim.create()
sim_client = SimClient(sim)
async with sim_and_client() as (sim, sim_client):
for i in range(0, 5):
await sim.farm_block()
await sim.farm_block(bytes32([0] * 32))
@ -143,6 +135,3 @@ class TestSimClient:
new_coin = next(x.coin for x in additions if x.coin.puzzle_hash == puzzle_hash)
coin_records = await sim_client.get_coin_records_by_parent_ids([spendable_coin.name()])
assert coin_records[0].coin.name() == new_coin.name()
finally:
await sim.close()

View File

@ -16,7 +16,6 @@ import pytest_asyncio
from _pytest.fixtures import SubRequest
# Set spawn after stdlib imports, but before other imports
from chia.clvm.spend_sim import SimClient, SpendSim
from chia.full_node.full_node_api import FullNodeAPI
from chia.protocols import full_node_protocol
from chia.server.server import ChiaServer
@ -701,14 +700,6 @@ async def timelord_service(bt):
yield _
@pytest_asyncio.fixture(scope="function")
async def setup_sim():
sim = await SpendSim.create()
sim_client = SimClient(sim)
await sim.farm_block()
return sim, sim_client
@pytest.fixture(scope="function")
def tmp_chia_root(tmp_path):
"""

View File

@ -7,7 +7,7 @@ import pytest
from blspy import G2Element
from chia_rs import Coin
from chia.clvm.spend_sim import SimClient, SpendSim
from chia.clvm.spend_sim import SimClient, SpendSim, sim_and_client
from chia.consensus.constants import ConsensusConstants
from chia.consensus.default_constants import DEFAULT_CONSTANTS
from chia.full_node.bitcoin_fee_estimator import BitcoinFeeEstimator
@ -24,6 +24,10 @@ the_puzzle_hash = bytes32(
bytes.fromhex("9dcf97a184f32623d11a73124ceb99a5709b083721e878a16d78f596718ba7b2")
) # Program.to(1)
NEW_DEFAULT_CONSTANTS: ConsensusConstants = DEFAULT_CONSTANTS.replace(
MAX_BLOCK_COST_CLVM=300000000, MEMPOOL_BLOCK_BUFFER=1
)
async def farm(
sim: SpendSim,
@ -47,11 +51,8 @@ def make_tx_sb(from_coin: Coin) -> SpendBundle:
async def init_test(
puzzle_hash: bytes32, spends_per_block: int
) -> Tuple[SpendSim, SimClient, BitcoinFeeEstimator, List[Coin], List[Coin]]:
defaults: ConsensusConstants = DEFAULT_CONSTANTS
sim = await SpendSim.create(defaults=defaults.replace(MAX_BLOCK_COST_CLVM=300000000, MEMPOOL_BLOCK_BUFFER=1))
cli = SimClient(sim)
sim: SpendSim, cli: SimClient, puzzle_hash: bytes32, spends_per_block: int
) -> Tuple[BitcoinFeeEstimator, List[Coin], List[Coin]]:
new_reward_coins = []
spend_coins = []
fee_coins = []
@ -73,75 +74,74 @@ async def init_test(
assert sim.block_records[1].reward_claims_incorporated[1].amount == 250000000000
estimator: BitcoinFeeEstimator = sim.mempool_manager.mempool.fee_estimator # type:ignore
return sim, cli, estimator, spend_coins, fee_coins # new_reward_coins
return estimator, spend_coins, fee_coins # new_reward_coins
@pytest.mark.asyncio
async def test_mempool_inclusion_filter_basic() -> None:
sim, cli, estimator, spend_coins, fee_coins = await init_test(the_puzzle_hash, 1)
assert len(sim.mempool_manager.mempool.spends) == 0
async with sim_and_client(defaults=NEW_DEFAULT_CONSTANTS, pass_prefarm=False) as (sim, cli):
estimator, spend_coins, fee_coins = await init_test(sim, cli, the_puzzle_hash, 1)
assert len(sim.mempool_manager.mempool.spends) == 0
spend_bundle: SpendBundle = make_tx_sb(spend_coins[0])
status, error = await cli.push_tx(spend_bundle)
assert len(sim.mempool_manager.mempool.spends) == 1
assert error is None
spend_bundle: SpendBundle = make_tx_sb(spend_coins[0])
status, error = await cli.push_tx(spend_bundle)
assert len(sim.mempool_manager.mempool.spends) == 1
assert error is None
mempool_item = sim.mempool_manager.get_mempool_item(spend_bundle.name())
assert mempool_item
mempool_item = sim.mempool_manager.get_mempool_item(spend_bundle.name())
assert mempool_item
def include_none(mm: MempoolManager, mi: MempoolItem) -> bool:
return False
def include_none(mm: MempoolManager, mi: MempoolItem) -> bool:
return False
def include_all(mm: MempoolManager, mi: MempoolItem) -> bool:
return True
def include_all(mm: MempoolManager, mi: MempoolItem) -> bool:
return True
additions, removals = await sim.farm_block(the_puzzle_hash, item_inclusion_filter=include_none)
assert len(sim.mempool_manager.mempool.spends) == 1
assert removals == []
additions, removals = await sim.farm_block(the_puzzle_hash, item_inclusion_filter=include_none)
assert len(sim.mempool_manager.mempool.spends) == 1
assert removals == []
additions, removals = await sim.farm_block(the_puzzle_hash, item_inclusion_filter=include_all)
assert len(sim.mempool_manager.mempool.spends) == 0
removal_ids = [c.name() for c in removals]
assert mempool_item.name not in removal_ids
await sim.close()
additions, removals = await sim.farm_block(the_puzzle_hash, item_inclusion_filter=include_all)
assert len(sim.mempool_manager.mempool.spends) == 0
removal_ids = [c.name() for c in removals]
assert mempool_item.name not in removal_ids
@pytest.mark.asyncio
async def test_mempoolitem_height_added(db_version: int) -> None:
sim, cli, estimator, spend_coins, fee_coins = await init_test(the_puzzle_hash, 1)
assert len(sim.mempool_manager.mempool.spends) == 0
async with sim_and_client(defaults=NEW_DEFAULT_CONSTANTS, pass_prefarm=False) as (sim, cli):
estimator, spend_coins, fee_coins = await init_test(sim, cli, the_puzzle_hash, 1)
assert len(sim.mempool_manager.mempool.spends) == 0
spend_bundle: SpendBundle = make_tx_sb(spend_coins[0])
spend_bundle: SpendBundle = make_tx_sb(spend_coins[0])
status, error = await cli.push_tx(spend_bundle)
assert len(sim.mempool_manager.mempool.spends) == 1
log.warning(f"{status, error} = cli.push_tx({spend_bundle.name()})")
status, error = await cli.push_tx(spend_bundle)
assert len(sim.mempool_manager.mempool.spends) == 1
log.warning(f"{status, error} = cli.push_tx({spend_bundle.name()})")
mempool_item = sim.mempool_manager.get_mempool_item(spend_bundle.name())
assert mempool_item
heights = {sim.get_height(): mempool_item.height_added_to_mempool}
def ignore_spend(mm: MempoolManager, mi: MempoolItem) -> bool:
mempool_item = sim.mempool_manager.get_mempool_item(spend_bundle.name())
assert mempool_item
return mi.name != mempool_item.name
heights = {sim.get_height(): mempool_item.height_added_to_mempool}
additions, removals = await sim.farm_block(the_puzzle_hash, item_inclusion_filter=ignore_spend)
removal_ids = [c.name() for c in removals]
assert mempool_item.name not in removal_ids
def ignore_spend(mm: MempoolManager, mi: MempoolItem) -> bool:
assert mempool_item
return mi.name != mempool_item.name
mempool_item2 = sim.mempool_manager.get_mempool_item(spend_bundle.name())
assert len(sim.mempool_manager.mempool.spends) == 1
assert mempool_item2
additions, removals = await sim.farm_block(the_puzzle_hash, item_inclusion_filter=ignore_spend)
removal_ids = [c.name() for c in removals]
assert mempool_item.name not in removal_ids
# This is the important check in this test: ensure height_added_to_mempool does not
# change when the mempool is rebuilt
assert mempool_item2.height_added_to_mempool == mempool_item2.height_added_to_mempool
mempool_item2 = sim.mempool_manager.get_mempool_item(spend_bundle.name())
assert len(sim.mempool_manager.mempool.spends) == 1
assert mempool_item2
# Now farm it into the next block
additions, removals = await sim.farm_block(the_puzzle_hash)
assert len(sim.mempool_manager.mempool.spends) == 0
assert len(removals) == 1
# This is the important check in this test: ensure height_added_to_mempool does not
# change when the mempool is rebuilt
assert mempool_item2.height_added_to_mempool == mempool_item2.height_added_to_mempool
log.warning(heights)
await sim.close()
# Now farm it into the next block
additions, removals = await sim.farm_block(the_puzzle_hash)
assert len(sim.mempool_manager.mempool.spends) == 0
assert len(removals) == 1
log.warning(heights)

View File

@ -6,7 +6,7 @@ import pytest
from blspy import AugSchemeMPL, G2Element, PrivateKey
from clvm.casts import int_to_bytes
from chia.clvm.spend_sim import SimClient, SpendSim
from chia.clvm.spend_sim import SimClient, SpendSim, sim_and_client
from chia.types.blockchain_format.coin import Coin
from chia.types.blockchain_format.program import Program
from chia.types.blockchain_format.sized_bytes import bytes32
@ -91,10 +91,8 @@ class TestCATLifecycle:
cost: Dict[str, int] = {}
@pytest.mark.asyncio()
async def test_cat_mod(self, setup_sim):
sim, sim_client = setup_sim
try:
async def test_cat_mod(self):
async with sim_and_client() as (sim, sim_client):
tail = Program.to([])
checker_solution = Program.to([])
cat_puzzle: Program = construct_cat_puzzle(CAT_MOD, tail.get_tree_hash(), acs)
@ -249,14 +247,9 @@ class TestCATLifecycle:
limitations_solutions=[checker_solution],
)
finally:
await sim.close()
@pytest.mark.asyncio()
async def test_complex_spend(self, setup_sim):
sim, sim_client = setup_sim
try:
async def test_complex_spend(self):
async with sim_and_client() as (sim, sim_client):
tail = Program.to([])
checker_solution = Program.to([])
cat_puzzle: Program = construct_cat_puzzle(CAT_MOD, tail.get_tree_hash(), acs)
@ -343,14 +336,10 @@ class TestCATLifecycle:
limitations_solutions=[checker_solution] * 4,
extra_deltas=[13, -21, 21, -13],
)
finally:
await sim.close()
@pytest.mark.asyncio()
async def test_genesis_by_id(self, setup_sim):
sim, sim_client = setup_sim
try:
async def test_genesis_by_id(self):
async with sim_and_client() as (sim, sim_client):
standard_acs = Program.to(1)
standard_acs_ph: bytes32 = standard_acs.get_tree_hash()
await sim.farm_block(standard_acs_ph)
@ -387,14 +376,9 @@ class TestCATLifecycle:
limitations_solutions=[checker_solution],
)
finally:
await sim.close()
@pytest.mark.asyncio()
async def test_genesis_by_puzhash(self, setup_sim):
sim, sim_client = setup_sim
try:
async def test_genesis_by_puzhash(self):
async with sim_and_client() as (sim, sim_client):
standard_acs = Program.to(1)
standard_acs_ph: bytes32 = standard_acs.get_tree_hash()
await sim.farm_block(standard_acs_ph)
@ -431,14 +415,9 @@ class TestCATLifecycle:
limitations_solutions=[checker_solution],
)
finally:
await sim.close()
@pytest.mark.asyncio()
async def test_everything_with_signature(self, setup_sim):
sim, sim_client = setup_sim
try:
async def test_everything_with_signature(self):
async with sim_and_client() as (sim, sim_client):
sk = PrivateKey.from_bytes(secret_exponent_for_index(1).to_bytes(32, "big"))
tail: Program = EverythingWithSig.construct([Program.to(sk.get_g1())])
checker_solution: Program = EverythingWithSig.solve([], {})
@ -541,14 +520,9 @@ class TestCATLifecycle:
additional_spends=[acs_bundle],
)
finally:
await sim.close()
@pytest.mark.asyncio()
async def test_delegated_tail(self, setup_sim):
sim, sim_client = setup_sim
try:
async def test_delegated_tail(self):
async with sim_and_client() as (sim, sim_client):
standard_acs = Program.to(1)
standard_acs_ph: bytes32 = standard_acs.get_tree_hash()
await sim.farm_block(standard_acs_ph)
@ -601,9 +575,6 @@ class TestCATLifecycle:
limitations_solutions=[checker_solution],
)
finally:
await sim.close()
def test_cost(self):
import json
import logging

View File

@ -5,6 +5,7 @@ from typing import Any, Dict, List, Optional
import pytest
from blspy import G2Element
from chia.clvm.spend_sim import sim_and_client
from chia.types.announcement import Announcement
from chia.types.blockchain_format.coin import Coin
from chia.types.blockchain_format.program import Program
@ -168,10 +169,8 @@ class TestOfferLifecycle:
cost: Dict[str, int] = {}
@pytest.mark.asyncio()
async def test_complex_offer(self, setup_sim):
sim, sim_client = setup_sim
try:
async def test_complex_offer(self):
async with sim_and_client() as (sim, sim_client):
coins_needed: Dict[Optional[str], List[int]] = {
None: [500, 400, 300],
"red": [250, 100],
@ -312,8 +311,6 @@ class TestOfferLifecycle:
assert result == (MempoolInclusionStatus.SUCCESS, None)
self.cost["complex offer"] = cost_of_spend_bundle(offer_bundle)
await sim.farm_block()
finally:
await sim.close()
def test_cost(self):
import json

View File

@ -5,7 +5,7 @@ from typing import Dict, List, Tuple
import pytest
from blspy import G2Element
from chia.clvm.spend_sim import SimClient, SpendSim
from chia.clvm.spend_sim import sim_and_client
from chia.types.blockchain_format.coin import Coin
from chia.types.blockchain_format.program import Program
from chia.types.blockchain_format.sized_bytes import bytes32
@ -39,9 +39,8 @@ NIL_PH = Program.to(None).get_tree_hash()
@pytest.mark.asyncio
async def test_graftroot(setup_sim: Tuple[SpendSim, SimClient]) -> None:
sim, sim_client = setup_sim
try:
async def test_graftroot() -> None:
async with sim_and_client() as (sim, sim_client):
# Create the coin we're testing
all_values: List[bytes32] = [bytes32([x] * 32) for x in range(0, 100)]
root, proofs = build_merkle_tree(all_values)
@ -138,5 +137,3 @@ async def test_graftroot(setup_sim: Tuple[SpendSim, SimClient]) -> None:
assert result == (MempoolInclusionStatus.FAILED, Err.GENERATOR_RUNTIME_ERROR)
with pytest.raises(ValueError, match="clvm raise"):
graftroot_puzzle.run(graftroot_spend.solution.to_program())
finally:
await sim.close()

View File

@ -1,12 +1,12 @@
from __future__ import annotations
import itertools
from typing import List, Tuple
from typing import List
import pytest
from blspy import G2Element
from chia.clvm.spend_sim import SimClient, SpendSim
from chia.clvm.spend_sim import sim_and_client
from chia.types.announcement import Announcement
from chia.types.blockchain_format.program import Program
from chia.types.blockchain_format.sized_bytes import bytes32
@ -28,10 +28,8 @@ ACS_PH = ACS.get_tree_hash()
@pytest.mark.asyncio()
@pytest.mark.parametrize("metadata_updater", ["default"])
async def test_state_layer(setup_sim: Tuple[SpendSim, SimClient], metadata_updater: str) -> None:
sim, sim_client = setup_sim
try:
async def test_state_layer(metadata_updater: str) -> None:
async with sim_and_client() as (sim, sim_client):
if metadata_updater == "default":
METADATA: Program = metadata_to_program(
{
@ -132,15 +130,11 @@ async def test_state_layer(setup_sim: Tuple[SpendSim, SimClient], metadata_updat
assert result == (MempoolInclusionStatus.SUCCESS, None)
await sim.farm_block()
state_layer_puzzle = create_nft_layer_puzzle_with_curry_params(metadata, METADATA_UPDATER_PUZZLE_HASH, ACS)
finally:
await sim.close()
@pytest.mark.asyncio()
async def test_ownership_layer(setup_sim: Tuple[SpendSim, SimClient]) -> None:
sim, sim_client = setup_sim
try:
async def test_ownership_layer() -> None:
async with sim_and_client() as (sim, sim_client):
TARGET_OWNER = bytes32([0] * 32)
TARGET_TP = Program.to([8]) # (x)
# (a (i 11 (q 4 19 (c 43 (q ()))) (q 8)) 1) or
@ -239,15 +233,11 @@ async def test_ownership_layer(setup_sim: Tuple[SpendSim, SimClient]) -> None:
TARGET_TP,
ACS,
).get_tree_hash()
finally:
await sim.close()
@pytest.mark.asyncio()
async def test_default_transfer_program(setup_sim: Tuple[SpendSim, SimClient]) -> None:
sim, sim_client = setup_sim
try:
async def test_default_transfer_program() -> None:
async with sim_and_client() as (sim, sim_client):
# Now make the ownership coin
FAKE_SINGLETON_MOD = Program.to([2, 5, 11]) # (a 5 11) | (mod (_ INNER_PUZ inner_sol) (a INNER_PUZ inner_sol))
FAKE_CAT_MOD = Program.to([2, 11, 23]) # (a 11 23) or (mod (_ _ INNER_PUZ inner_sol) (a INNER_PUZ inner_sol))
@ -363,5 +353,3 @@ async def test_default_transfer_program(setup_sim: Tuple[SpendSim, SimClient]) -
result = await sim_client.push_tx(empty_bundle)
assert result == (MempoolInclusionStatus.SUCCESS, None)
await sim.farm_block()
finally:
await sim.close()