Simulator Final touches (#12703)

* add new rpc and add tests.

I added, get_all_blocks rpc, and expanded the get_all_puzzle_hashes rpc.

* lint pass

* update plot directory for simulator

* return peak height on new block farm

* expand test coverage

* fix plot directory craziness

* simulator sync override

* Add new simulator networktype

* add logging to simulator

* Fix race condition, mempool and confirm

* add reorg seed

* farm block if block 0 sim & fix unsynced for sim

* lint

* oops i forgot to lint

* switch to detecting simulator in network type

also remove inline genesis generation

* fix config loading to use correct variable

Co-authored-by: Mariano Sorgente <sorgente711@gmail.com>

chia/full_node/full_node.py

@@ -669,6 +669,8 @@ class FullNode:
                 await self.server.send_to_specific([msg], peer.peer_node_id)
 
     async def synced(self) -> bool:
+        if "simulator" in str(self.config.get("selected_network")):
+            return True  # sim is always synced because it has no peers
         curr: Optional[BlockRecord] = self.blockchain.get_peak()
         if curr is None:
             return False

chia/rpc/full_node_rpc_api.py

@@ -186,7 +186,9 @@ class FullNodeRpcApi:
             mempool_min_fee_5m = 0
             mempool_max_total_cost = 0
         if self.service.server is not None:
-            is_connected = len(self.service.server.get_full_node_connections()) > 0
+            is_connected = len(self.service.server.get_full_node_connections()) > 0 or "simulator" in str(
+                self.service.config.get("selected_network")
+            )
         else:
             is_connected = False
         synced = await self.service.synced() and is_connected

chia/simulator/block_tools.py

@@ -127,9 +127,6 @@ test_constants = DEFAULT_CONSTANTS.replace(
 )
 
 
-log = logging.getLogger(__name__)
-
-
 class BlockTools:
     """
     Tools to generate blocks for testing.
@@ -149,6 +146,7 @@ class BlockTools:
         config_overrides: Optional[Dict] = None,
         automated_testing: bool = True,
         plot_dir: str = "test-plots",
+        log: logging.Logger = logging.getLogger(__name__),
     ):
 
         self._block_cache_header = bytes32([0] * 32)
@@ -159,6 +157,7 @@ class BlockTools:
             root_path = Path(self._tempdir.name)
 
         self.root_path = root_path
+        self.log = log
         self.local_keychain = keychain
         self._block_time_residual = 0.0
         self.local_sk_cache: Dict[bytes32, Tuple[PrivateKey, Any]] = {}
@@ -238,12 +237,12 @@ class BlockTools:
 
     async def setup_keys(self, fingerprint: Optional[int] = None, reward_ph: Optional[bytes32] = None):
         if self.local_keychain:
-            keychain_proxy: Optional[KeychainProxy] = wrap_local_keychain(self.local_keychain, log=log)
+            keychain_proxy: Optional[KeychainProxy] = wrap_local_keychain(self.local_keychain, log=self.log)
         elif not self.automated_testing and fingerprint is not None:
-            keychain_proxy = await connect_to_keychain_and_validate(self.root_path, log)
+            keychain_proxy = await connect_to_keychain_and_validate(self.root_path, self.log)
         else:  # if we are automated testing or if we don't have a fingerprint.
             keychain_proxy = await connect_to_keychain_and_validate(
-                self.root_path, log, user="testing-1.8.0", service="chia-testing-1.8.0"
+                self.root_path, self.log, user="testing-1.8.0", service="chia-testing-1.8.0"
             )
         assert keychain_proxy is not None
         if fingerprint is None:  # if we are not specifying an existing key
@@ -520,7 +519,7 @@ class BlockTools:
                 skip_slots=skip_slots,
                 timestamp=(uint64(int(time.time())) if genesis_timestamp is None else genesis_timestamp),
             )
-            log.info(f"Created block 0 iters: {genesis.total_iters}")
+            self.log.info(f"Created block 0 iters: {genesis.total_iters}")
             num_empty_slots_added = skip_slots
             block_list = [genesis]
             num_blocks -= 1
@@ -718,7 +717,9 @@ class BlockTools:
                         blocks_added_this_sub_slot += 1
 
                         blocks[full_block.header_hash] = block_record
-                        log.info(f"Created block {block_record.height} ove=False, iters " f"{block_record.total_iters}")
+                        self.log.info(
+                            f"Created block {block_record.height} ove=False, iters " f"{block_record.total_iters}"
+                        )
                         height_to_hash[uint32(full_block.height)] = full_block.header_hash
                         latest_block = blocks[full_block.header_hash]
                         finished_sub_slots_at_ip = []
@@ -796,7 +797,7 @@ class BlockTools:
                 new_sub_slot_iters: Optional[uint64] = sub_epoch_summary.new_sub_slot_iters
                 new_difficulty: Optional[uint64] = sub_epoch_summary.new_difficulty
 
-                log.info(f"Sub epoch summary: {sub_epoch_summary}")
+                self.log.info(f"Sub epoch summary: {sub_epoch_summary}")
             else:
                 ses_hash = None
                 new_sub_slot_iters = None
@@ -870,7 +871,7 @@ class BlockTools:
                 additions = transaction_data.additions()
                 removals = transaction_data.removals()
             sub_slots_finished += 1
-            log.info(
+            self.log.info(
                 f"Sub slot finished. blocks included: {blocks_added_this_sub_slot} blocks_per_slot: "
                 f"{(len(block_list) - initial_block_list_len)/sub_slots_finished}"
             )
@@ -995,7 +996,9 @@ class BlockTools:
                                 previous_generator = compressor_arg
 
                         blocks_added_this_sub_slot += 1
-                        log.info(f"Created block {block_record.height } ov=True, iters " f"{block_record.total_iters}")
+                        self.log.info(
+                            f"Created block {block_record.height } ov=True, iters " f"{block_record.total_iters}"
+                        )
                         num_blocks -= 1
 
                         blocks[full_block.header_hash] = block_record
@@ -1439,7 +1442,10 @@ def get_challenges(
 
 
 def get_plot_dir(plot_dir_name: str = "test-plots", automated_testing: bool = True) -> Path:
-    cache_path = DEFAULT_ROOT_PATH.parent.joinpath(plot_dir_name)
+    root_dir = DEFAULT_ROOT_PATH.parent
+    if not automated_testing:  # make sure we don't accidentally stack directories.
+        root_dir = root_dir.parent if root_dir.parts[-1] == plot_dir_name.split("/")[0] else root_dir
+    cache_path = root_dir.joinpath(plot_dir_name)
 
     ci = os.environ.get("CI")
     if ci is not None and not cache_path.exists() and automated_testing:

chia/simulator/full_node_simulator.py

@@ -1,6 +1,6 @@
 import asyncio
 import time
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Tuple
 
 from chia.consensus.block_record import BlockRecord
 from chia.consensus.multiprocess_validation import PreValidationResult
@@ -100,15 +100,16 @@ class FullNodeSimulator(FullNodeAPI):
         # reload mempool
         await self.full_node.mempool_manager.new_peak(block_record, None)
 
-    async def get_all_puzzle_hashes(self) -> Dict[bytes32, uint128]:
-        # puzzlehash, total_amount
-        ph_total_amount: Dict[bytes32, uint128] = {}
+    async def get_all_puzzle_hashes(self) -> Dict[bytes32, Tuple[uint128, int]]:
+        # puzzle_hash, (total_amount, num_transactions)
+        ph_total_amount: Dict[bytes32, Tuple[uint128, int]] = {}
         all_non_spent_coins: List[CoinRecord] = await self.get_all_coins(GetAllCoinsProtocol(False))
         for cr in all_non_spent_coins:
             if cr.coin.puzzle_hash not in ph_total_amount:
-                ph_total_amount[cr.coin.puzzle_hash] = uint128(cr.coin.amount)
+                ph_total_amount[cr.coin.puzzle_hash] = (uint128(cr.coin.amount), 1)
             else:
-                ph_total_amount[cr.coin.puzzle_hash] = uint128(cr.coin.amount + ph_total_amount[cr.coin.puzzle_hash])
+                dict_value: Tuple[uint128, int] = ph_total_amount[cr.coin.puzzle_hash]
+                ph_total_amount[cr.coin.puzzle_hash] = (uint128(cr.coin.amount + dict_value[0]), dict_value[1] + 1)
         return ph_total_amount
 
     @api_request
@@ -215,6 +216,9 @@ class FullNodeSimulator(FullNodeAPI):
         new_index = request.new_index
         old_index = request.old_index
         coinbase_ph = request.puzzle_hash
+        seed = request.seed
+        if seed is None:
+            seed = bytes32(32 * b"1")
 
         current_blocks = await self.get_all_full_blocks()
         block_count = new_index - old_index
@@ -226,7 +230,7 @@ class FullNodeSimulator(FullNodeAPI):
             block_list_input=current_blocks[: old_index + 1],
             force_overflow=True,
             guarantee_transaction_block=True,
-            seed=32 * b"1",
+            seed=seed,
         )
 
         for block in more_blocks:

chia/simulator/simulator_full_node_rpc_api.py

@@ -1,9 +1,12 @@
+from secrets import token_bytes
 from typing import Dict, List
 
 from chia.rpc.full_node_rpc_api import FullNodeRpcApi
 from chia.rpc.rpc_server import Endpoint, EndpointResult
 from chia.simulator.simulator_protocol import FarmNewBlockProtocol, GetAllCoinsProtocol, ReorgProtocol
+from chia.types.blockchain_format.sized_bytes import bytes32
 from chia.types.coin_record import CoinRecord
+from chia.types.full_block import FullBlock
 from chia.util.bech32m import decode_puzzle_hash
 from chia.util.ints import uint32
 
@@ -11,6 +14,7 @@ from chia.util.ints import uint32
 class SimulatorFullNodeRpcApi(FullNodeRpcApi):
     def get_routes(self) -> Dict[str, Endpoint]:
         routes = super().get_routes()
+        routes["/get_all_blocks"] = self.get_all_blocks
         routes["/farm_block"] = self.farm_block
         routes["/set_auto_farming"] = self.set_auto_farming
         routes["/get_auto_farming"] = self.get_auto_farming
@@ -21,19 +25,24 @@ class SimulatorFullNodeRpcApi(FullNodeRpcApi):
         routes["/reorg_blocks"] = self.reorg_blocks
         return routes
 
+    async def get_all_blocks(self, _request: Dict[str, object]) -> EndpointResult:
+        all_blocks: List[FullBlock] = await self.service.server.api.get_all_full_blocks()
+        return {"blocks": [block.to_json_dict() for block in all_blocks]}
+
     async def farm_block(self, _request: Dict[str, object]) -> EndpointResult:
         request_address = str(_request["address"])
         guarantee_tx_block = bool(_request.get("guarantee_tx_block", False))
         blocks = int(str(_request.get("blocks", 1)))  # mypy made me do this
         ph = decode_puzzle_hash(request_address)
         req = FarmNewBlockProtocol(ph)
+        cur_height = self.service.blockchain.get_peak_height()
         if guarantee_tx_block:
             for i in range(blocks):  # these can only be tx blocks
                 await self.service.server.api.farm_new_transaction_block(req)
         else:
             for i in range(blocks):  # these can either be full blocks or tx blocks
                 await self.service.server.api.farm_new_block(req)
-        return {}
+        return {"new_peak_height": (cur_height if cur_height is not None else 0) + blocks}
 
     async def set_auto_farming(self, _request: Dict[str, object]) -> EndpointResult:
         auto_farm = bool(_request["auto_farm"])
@@ -53,7 +62,9 @@ class SimulatorFullNodeRpcApi(FullNodeRpcApi):
 
     async def get_all_puzzle_hashes(self, _request: Dict[str, object]) -> EndpointResult:
         result = await self.service.server.api.get_all_puzzle_hashes()
-        return {"puzzle_hashes": {puzzle_hash.hex(): amount for (puzzle_hash, amount) in result.items()}}
+        return {
+            "puzzle_hashes": {puzzle_hash.hex(): (amount, num_tx) for (puzzle_hash, (amount, num_tx)) in result.items()}
+        }
 
     async def revert_blocks(self, _request: Dict[str, object]) -> EndpointResult:
         blocks = int(str(_request.get("num_of_blocks", 1)))  # number of blocks to revert
@@ -70,12 +81,16 @@ class SimulatorFullNodeRpcApi(FullNodeRpcApi):
         fork_blocks = int(str(_request.get("num_of_blocks_to_rev", 1)))  # number of blocks to go back
         new_blocks = int(str(_request.get("num_of_new_blocks", 1)))  # how many extra blocks should we add
         all_blocks = bool(_request.get("revert_all_blocks", False))  # fork all blocks
+        use_random_seed = bool(_request.get("random_seed", True))  # randomize the seed to differentiate reorgs
+        random_seed = bytes32(token_bytes(32)) if use_random_seed else None
         cur_height = self.service.blockchain.get_peak_height()
         if cur_height is None:
             raise ValueError("No blocks to revert")
         fork_height = (cur_height - fork_blocks) if not all_blocks else 1
         new_height = cur_height + new_blocks  # any number works as long as its not 0
         assert fork_height >= 1 and new_height - 1 >= cur_height
-        request = ReorgProtocol(uint32(fork_height), uint32(new_height), self.service.server.api.bt.farmer_ph)
+        request = ReorgProtocol(
+            uint32(fork_height), uint32(new_height), self.service.server.api.bt.farmer_ph, random_seed
+        )
         await self.service.server.api.reorg_from_index_to_new_index(request)
         return {"new_peak_height": new_height}

chia/simulator/simulator_full_node_rpc_client.py

@@ -1,25 +1,29 @@
-from typing import Dict, List
+from typing import Dict, List, Tuple
 
 from chia.rpc.full_node_rpc_client import FullNodeRpcClient
 from chia.types.blockchain_format.sized_bytes import bytes32
 from chia.types.coin_record import CoinRecord
+from chia.types.full_block import FullBlock
 from chia.util.bech32m import encode_puzzle_hash
 from chia.util.byte_types import hexstr_to_bytes
 from chia.util.ints import uint128
 
 
 class SimulatorFullNodeRpcClient(FullNodeRpcClient):
-    async def farm_block(self, target_ph: bytes32, number_of_blocks: int = 1, guarantee_tx_block: bool = False) -> None:
+    async def get_all_blocks(self) -> List[FullBlock]:
+        json_blocks = (await self.fetch("get_all_blocks", {}))["blocks"]
+        return [FullBlock.from_json_dict(block) for block in json_blocks]
+
+    async def farm_block(self, target_ph: bytes32, number_of_blocks: int = 1, guarantee_tx_block: bool = False) -> int:
         address = encode_puzzle_hash(target_ph, "txch")
-        await self.fetch(
-            "farm_block", {"address": address, "blocks": number_of_blocks, "guarantee_tx_block": guarantee_tx_block}
-        )
+        request_args = {"address": address, "blocks": number_of_blocks, "guarantee_tx_block": guarantee_tx_block}
+        new_height: int = (await self.fetch("farm_block", request_args))["new_peak_height"]
+        return new_height
 
     async def set_auto_farming(self, set_auto_farming: bool) -> bool:
-        result = await self.fetch("set_auto_farming", {"auto_farm": set_auto_farming})
-        result = result["auto_farm_enabled"]
+        result: bool = (await self.fetch("set_auto_farming", {"auto_farm": set_auto_farming}))["auto_farm_enabled"]
         assert result == set_auto_farming
-        return bool(result)
+        return result
 
     async def get_auto_farming(self) -> bool:
         result = await self.fetch("get_auto_farming", {})
@@ -33,20 +37,25 @@ class SimulatorFullNodeRpcClient(FullNodeRpcClient):
         json_result = await self.fetch("get_all_coins", {"include_spent_coins": include_spent_coins})
         return [CoinRecord.from_json_dict(coin_records) for coin_records in json_result["coin_records"]]
 
-    async def get_all_puzzle_hashes(self) -> Dict[bytes32, uint128]:
+    async def get_all_puzzle_hashes(self) -> Dict[bytes32, Tuple[uint128, int]]:
         str_result = (await self.fetch("get_all_puzzle_hashes", {}))["puzzle_hashes"]
-        return {bytes32.from_hexstr(ph): uint128(amount) for (ph, amount) in str_result.items()}
+        return {bytes32.from_hexstr(ph): (uint128(amount), num_tx) for (ph, (amount, num_tx)) in str_result.items()}
 
     async def revert_blocks(self, num_of_blocks_to_delete: int = 1, delete_all_blocks: bool = False) -> int:
         request = {"delete_all_blocks": delete_all_blocks, "num_of_blocks": num_of_blocks_to_delete}
         return int((await self.fetch("revert_blocks", request))["new_peak_height"])
 
     async def reorg_blocks(
-        self, num_of_blocks_to_revert: int = 1, num_of_new_blocks: int = 1, revert_all_blocks: bool = False
+        self,
+        num_of_blocks_to_revert: int = 1,
+        num_of_new_blocks: int = 1,
+        revert_all_blocks: bool = False,
+        random_seed: bool = True,
     ) -> int:
         request = {
             "revert_all_blocks": revert_all_blocks,
             "num_of_blocks_to_rev": num_of_blocks_to_revert,
             "num_of_new_blocks": num_of_new_blocks,
+            "random_seed": random_seed,
         }
         return int((await self.fetch("reorg_blocks", request))["new_peak_height"])

chia/simulator/simulator_protocol.py

@@ -1,4 +1,5 @@
 from dataclasses import dataclass
+from typing import Optional
 
 from chia.types.blockchain_format.sized_bytes import bytes32
 from chia.util.ints import uint32
@@ -17,6 +18,7 @@ class ReorgProtocol(Streamable):
     old_index: uint32
     new_index: uint32
     puzzle_hash: bytes32
+    seed: Optional[bytes32]
 
 
 @streamable

chia/simulator/start_simulator.py

@@ -1,3 +1,4 @@
+import logging
 import sys
 from multiprocessing import freeze_support
 from pathlib import Path
@@ -9,7 +10,8 @@ from chia.server.start_service import Service, async_run
 from chia.simulator.simulator_full_node_rpc_api import SimulatorFullNodeRpcApi
 from chia.types.blockchain_format.sized_bytes import bytes32
 from chia.util.bech32m import decode_puzzle_hash
-from chia.util.config import load_config_cli, override_config
+from chia.util.chia_logging import initialize_logging
+from chia.util.config import load_config_cli, override_config, load_config
 from chia.util.default_root import DEFAULT_ROOT_PATH
 from chia.util.path import path_from_root
 from chia.simulator.block_tools import BlockTools, test_constants
@@ -20,6 +22,7 @@ from chia.simulator.full_node_simulator import FullNodeSimulator
 "".encode("idna")
 
 SERVICE_NAME = "full_node"
+log = logging.getLogger(__name__)
 PLOTS = 3  # 3 plots should be enough
 PLOT_SIZE = 19  # anything under k19 is a bit buggy
 
@@ -39,7 +42,6 @@ def create_full_node_simulator_service(
         config=service_config,
         root_path=root_path,
         consensus_constants=constants,
-        name=SERVICE_NAME,
     )
 
     peer_api = FullNodeSimulator(node, bt, config)
@@ -62,15 +64,17 @@ def create_full_node_simulator_service(
 
 
 async def async_main(test_mode: bool = False, automated_testing: bool = False, root_path: Path = DEFAULT_ROOT_PATH):
-    # We always use a real keychain for the new simulator.
-    config = load_config_cli(root_path, "config.yaml")
-    service_config = config[SERVICE_NAME]
+    # Same as full node, but the root_path is defined above
+    config = load_config(root_path, "config.yaml")
+    service_config = load_config_cli(root_path, "config.yaml", SERVICE_NAME)
+    config[SERVICE_NAME] = service_config
+    # THIS IS Simulator specific.
     fingerprint: Optional[int] = None
     farming_puzzle_hash: Optional[bytes32] = None
-    plot_dir: str = "simulator-plots"
+    plot_dir: str = "simulator/plots"
     if "simulator" in config:
         overrides = {}
-        plot_dir = config["simulator"].get("plot_directory", "simulator-plots")
+        plot_dir = config["simulator"].get("plot_directory", "simulator/plots")
         if config["simulator"]["key_fingerprint"] is not None:
             fingerprint = int(config["simulator"]["key_fingerprint"])
         if config["simulator"]["farming_address"] is not None:
@@ -94,9 +98,16 @@ async def async_main(test_mode: bool = False, automated_testing: bool = False, r
     )
     await bt.setup_keys(fingerprint=fingerprint, reward_ph=farming_puzzle_hash)
     await bt.setup_plots(num_og_plots=PLOTS, num_pool_plots=0, num_non_keychain_plots=0, plot_size=PLOT_SIZE)
+    # Everything after this is not simulator specific, excluding the if test_mode.
+    initialize_logging(
+        service_name=SERVICE_NAME,
+        logging_config=service_config["logging"],
+        root_path=DEFAULT_ROOT_PATH,
+    )
     service = create_full_node_simulator_service(root_path, override_config(config, overrides), bt)
     if test_mode:
         return service
+    await service.setup_process_global_state()
     await service.run()
     return 0
 

chia/util/initial-config.yaml

@@ -565,8 +565,8 @@ simulator:
   farming_address:
 
   # the directory used to save plots
-  # for the plot directory below the final directory will be: ~/.chia/simulator-plots
-  plot_directory: simulator-plots
+  # for the plot directory below the final directory will be: ~/.chia/simulator/plots
+  plot_directory: simulator/plots
 
   # Should we use real time in the simulated chain?
   # most tests don't need this, however it is pretty important when writing ChiaLisp

chia/wallet/wallet_node_api.py

@@ -76,30 +76,33 @@ class WalletNodeAPI:
         This is an ack for our previous SendTransaction call. This removes the transaction from
         the send queue if we have sent it to enough nodes.
         """
-        assert peer.peer_node_id is not None
-        name = peer.peer_node_id.hex()
-        status = MempoolInclusionStatus(ack.status)
-        try:
-            wallet_state_manager = self.wallet_node.wallet_state_manager
-        except RuntimeError as e:
-            if "not assigned" in str(e):
-                return None
-            raise
+        async with self.wallet_node.wallet_state_manager.lock:
+            assert peer.peer_node_id is not None
+            name = peer.peer_node_id.hex()
+            status = MempoolInclusionStatus(ack.status)
+            try:
+                wallet_state_manager = self.wallet_node.wallet_state_manager
+            except RuntimeError as e:
+                if "not assigned" in str(e):
+                    return None
+                raise
 
-        if status == MempoolInclusionStatus.SUCCESS:
-            self.wallet_node.log.info(f"SpendBundle has been received and accepted to mempool by the FullNode. {ack}")
-        elif status == MempoolInclusionStatus.PENDING:
-            self.wallet_node.log.info(f"SpendBundle has been received (and is pending) by the FullNode. {ack}")
-        else:
-            if not self.wallet_node.is_trusted(peer) and ack.error == Err.NO_TRANSACTIONS_WHILE_SYNCING.name:
-                self.wallet_node.log.info(f"Peer {peer.get_peer_info()} is not synced, closing connection")
-                await peer.close()
-                return
-            self.wallet_node.log.warning(f"SpendBundle has been rejected by the FullNode. {ack}")
-        if ack.error is not None:
-            await wallet_state_manager.remove_from_queue(ack.txid, name, status, Err[ack.error])
-        else:
-            await wallet_state_manager.remove_from_queue(ack.txid, name, status, None)
+            if status == MempoolInclusionStatus.SUCCESS:
+                self.wallet_node.log.info(
+                    f"SpendBundle has been received and accepted to mempool by the FullNode. {ack}"
+                )
+            elif status == MempoolInclusionStatus.PENDING:
+                self.wallet_node.log.info(f"SpendBundle has been received (and is pending) by the FullNode. {ack}")
+            else:
+                if not self.wallet_node.is_trusted(peer) and ack.error == Err.NO_TRANSACTIONS_WHILE_SYNCING.name:
+                    self.wallet_node.log.info(f"Peer {peer.get_peer_info()} is not synced, closing connection")
+                    await peer.close()
+                    return
+                self.wallet_node.log.warning(f"SpendBundle has been rejected by the FullNode. {ack}")
+            if ack.error is not None:
+                await wallet_state_manager.remove_from_queue(ack.txid, name, status, Err[ack.error])
+            else:
+                await wallet_state_manager.remove_from_queue(ack.txid, name, status, None)
 
     @peer_required
     @api_request

chia/wallet/wallet_state_manager.py

@@ -514,6 +514,9 @@ class WalletStateManager:
         if latest is None:
             return False
 
+        if "simulator" in self.config.get("selected_network"):
+            return True  # sim is always synced if we have a genesis block.
+
         if latest.height - await self.blockchain.get_finished_sync_up_to() > 1:
             return False
 

tests/core/test_full_node_rpc.py

@@ -270,7 +270,7 @@ class TestRpc:
             blocks: List[FullBlock] = await client.get_blocks(0, 5)
             assert len(blocks) == 5
 
-            await full_node_api_1.reorg_from_index_to_new_index(ReorgProtocol(2, 55, bytes([0x2] * 32)))
+            await full_node_api_1.reorg_from_index_to_new_index(ReorgProtocol(2, 55, bytes([0x2] * 32), None))
             new_blocks_0: List[FullBlock] = await client.get_blocks(0, 5)
             assert len(new_blocks_0) == 7
 

tests/pools/test_pool_rpc.py

@@ -373,7 +373,7 @@ class TestPoolWalletRpc:
         assert len(await wallet_node_0.wallet_state_manager.tx_store.get_unconfirmed_for_wallet(2)) == 0
         assert len(await wallet_node_0.wallet_state_manager.tx_store.get_unconfirmed_for_wallet(3)) == 0
         # Doing a reorg reverts and removes the pool wallets
-        await full_node_api.reorg_from_index_to_new_index(ReorgProtocol(uint32(0), uint32(20), our_ph_2))
+        await full_node_api.reorg_from_index_to_new_index(ReorgProtocol(uint32(0), uint32(20), our_ph_2, None))
         await time_out_assert(30, wallet_is_synced, True, wallet_node_0, full_node_api)
         summaries_response = await client.get_wallets()
         assert len(summaries_response) == 1

tests/simulation/test_simulation.py

@@ -172,7 +172,7 @@ class TestSimulation:
         assert len(spent_and_non_spent_coins) == 12
         # try reorg, then check again.
         # revert to height 2, then go to height 6, so that we don't include the transaction we made.
-        await full_node_api.reorg_from_index_to_new_index(ReorgProtocol(uint32(2), uint32(6), ph))
+        await full_node_api.reorg_from_index_to_new_index(ReorgProtocol(uint32(2), uint32(6), ph, None))
         reorg_non_spent_coins = await full_node_api.get_all_coins(GetAllCoinsProtocol(False))
         reorg_spent_and_non_spent_coins = await full_node_api.get_all_coins(GetAllCoinsProtocol(True))
         assert len(reorg_non_spent_coins) == 12 and len(reorg_spent_and_non_spent_coins) == 12

tests/simulation/test_start_simulator.py

@@ -85,7 +85,7 @@ def create_config(chia_root: Path, fingerprint: int) -> Dict[str, Any]:
     # simulator overrides
     config["simulator"]["key_fingerprint"] = fingerprint
     config["simulator"]["farming_address"] = encode_puzzle_hash(get_puzzle_hash_from_key(fingerprint), "txch")
-    config["simulator"]["plot_directory"] = "test-simulator-plots"
+    config["simulator"]["plot_directory"] = "test-simulator/plots"
     # save config
     save_config(chia_root, "config.yaml", config)
     return config
@@ -165,13 +165,14 @@ class TestStartSimulator:
         await time_out_assert(10, get_num_coins_for_ph, 2, simulator_rpc_client, ph_1)
         # test both block RPC's
         await simulator_rpc_client.farm_block(ph_2)
-        await simulator_rpc_client.farm_block(ph_2, guarantee_tx_block=True)
+        new_height = await simulator_rpc_client.farm_block(ph_2, guarantee_tx_block=True)
         # check if farming reward was received correctly & if block was created
-        await time_out_assert(10, simulator.full_node.blockchain.get_peak_height, 4)
+        await time_out_assert(10, simulator.full_node.blockchain.get_peak_height, new_height)
         await time_out_assert(10, get_num_coins_for_ph, 2, simulator_rpc_client, ph_2)
         # test balance rpc
         ph_amount = await simulator_rpc_client.get_all_puzzle_hashes()
-        assert ph_amount[ph_2] == 2000000000000
+        assert ph_amount[ph_2][0] == 2000000000000
+        assert ph_amount[ph_2][1] == 2
         # test all coins rpc.
         coin_records = await simulator_rpc_client.get_all_coins()
         ph_2_total = 0
@@ -184,7 +185,7 @@ class TestStartSimulator:
         assert ph_2_total == 2000000000000 and ph_1_total == 4000000000000
         # block rpc tests.
         # test reorg
-        old_blocks = await simulator.get_all_full_blocks()  # len should be 4
+        old_blocks = await simulator_rpc_client.get_all_blocks()  # len should be 4
         await simulator_rpc_client.reorg_blocks(2)  # fork point 2 blocks, now height is 5
         await time_out_assert(10, simulator.full_node.blockchain.get_peak_height, 5)
         # now validate that the blocks don't match

tests/wallet/cat_wallet/test_cat_wallet.py

@@ -99,7 +99,9 @@ class TestCATWallet:
         assert await cat_wallet.lineage_store.get_all_lineage_proofs() == all_lineage
 
         height = full_node_api.full_node.blockchain.get_peak_height()
-        await full_node_api.reorg_from_index_to_new_index(ReorgProtocol(height - num_blocks - 1, height + 1, 32 * b"1"))
+        await full_node_api.reorg_from_index_to_new_index(
+            ReorgProtocol(height - num_blocks - 1, height + 1, 32 * b"1", None)
+        )
         await time_out_assert(20, cat_wallet.get_confirmed_balance, 0)
 
     @pytest.mark.asyncio
@@ -243,7 +245,7 @@ class TestCATWallet:
         await time_out_assert(20, cat_wallet.get_unconfirmed_balance, 55)
 
         height = full_node_api.full_node.blockchain.get_peak_height()
-        await full_node_api.reorg_from_index_to_new_index(ReorgProtocol(height - 1, height + 1, 32 * b"1"))
+        await full_node_api.reorg_from_index_to_new_index(ReorgProtocol(height - 1, height + 1, 32 * b"1", None))
         await time_out_assert(20, cat_wallet.get_confirmed_balance, 40)
 
     @pytest.mark.parametrize(

tests/wallet/simple_sync/test_simple_sync_protocol.py

@@ -383,7 +383,7 @@ class TestSimpleSyncProtocol:
         coin_records = await full_node_api.full_node.coin_store.get_coin_records_by_puzzle_hash(True, puzzle_hash)
         assert len(coin_records) > 0
         fork_height = expected_height - num_blocks - 5
-        req = ReorgProtocol(fork_height, expected_height + 5, zero_ph)
+        req = ReorgProtocol(fork_height, expected_height + 5, zero_ph, None)
         await full_node_api.reorg_from_index_to_new_index(req)
 
         coin_records = await full_node_api.full_node.coin_store.get_coin_records_by_puzzle_hash(True, puzzle_hash)
@@ -461,7 +461,7 @@ class TestSimpleSyncProtocol:
             assert msg_response is not None
 
         fork_height = expected_height - num_blocks - 5
-        req = ReorgProtocol(fork_height, expected_height + 5, zero_ph)
+        req = ReorgProtocol(fork_height, expected_height + 5, zero_ph, None)
         await full_node_api.reorg_from_index_to_new_index(req)
 
         coin_records = await full_node_api.full_node.coin_store.get_coin_records_by_puzzle_hash(True, puzzle_hash)

tests/wallet/test_wallet.py

@@ -185,7 +185,7 @@ class TestWalletSimulator:
         await time_out_assert(25, wallet.get_confirmed_balance, funds)
 
         await full_node_api.reorg_from_index_to_new_index(
-            ReorgProtocol(uint32(2), uint32(num_blocks + 6), bytes32(32 * b"0"))
+            ReorgProtocol(uint32(2), uint32(num_blocks + 6), bytes32(32 * b"0"), None)
         )
 
         funds = sum(
@@ -735,7 +735,7 @@ class TestWalletSimulator:
 
         # Perform a reorg, which will revert the transaction in the full node and wallet, and cause wallet to resubmit
         await full_node_api.reorg_from_index_to_new_index(
-            ReorgProtocol(uint32(peak_height - 3), uint32(peak_height + 3), bytes32(32 * b"0"))
+            ReorgProtocol(uint32(peak_height - 3), uint32(peak_height + 3), bytes32(32 * b"0"), None)
         )
 
         funds = sum(