context manager for BlockTools

benchmarks/blockchains.py

@@ -33,25 +33,25 @@ def enable_profiler(profile: bool, name: str) -> Iterator[None]:
 
 async def run_test_chain_benchmark() -> None:
     with TempKeyring() as keychain:
-        bt = await create_block_tools_async(constants=test_constants, keychain=keychain)
-        with enable_profiler(True, "load-test-chain"):
-            start = time.monotonic()
-            for version in ["", "_hardfork"]:
-                for count, name in [
-                    (400, "test_blocks_400_rc5"),
-                    (1000, "test_blocks_1000_rc5"),
-                    (1000, "pre_genesis_empty_slots_1000_blocksrc5"),
-                    (1500, "test_blocks_1500_rc5"),
-                    (10000, "test_blocks_10000_rc5"),
-                    (758 + 320, "test_blocks_long_reorg_rc5"),
-                    (2000, "test_blocks_2000_compact_rc5"),
-                    (10000, "test_blocks_10000_compact_rc5"),
-                ]:
-                    persistent_blocks(count, f"{name}{version}.db", bt, seed=b"100")
-            end = time.monotonic()
-        KeyringWrapper.cleanup_shared_instance()
+        async with create_block_tools_async(constants=test_constants, keychain=keychain) as bt:
+            with enable_profiler(True, "load-test-chain"):
+                start = time.monotonic()
+                for version in ["", "_hardfork"]:
+                    for count, name in [
+                        (400, "test_blocks_400_rc5"),
+                        (1000, "test_blocks_1000_rc5"),
+                        (1000, "pre_genesis_empty_slots_1000_blocksrc5"),
+                        (1500, "test_blocks_1500_rc5"),
+                        (10000, "test_blocks_10000_rc5"),
+                        (758 + 320, "test_blocks_long_reorg_rc5"),
+                        (2000, "test_blocks_2000_compact_rc5"),
+                        (10000, "test_blocks_10000_compact_rc5"),
+                    ]:
+                        persistent_blocks(count, f"{name}{version}.db", bt, seed=b"100")
+                end = time.monotonic()
+            KeyringWrapper.cleanup_shared_instance()
 
-        print(f"time to load test chains: {end - start:.2f}s")
+            print(f"time to load test chains: {end - start:.2f}s")
 
 
 if __name__ == "__main__":

chia/_tests/blockchain/test_blockchain.py

@@ -618,14 +618,16 @@ async def test_genesis_no_icc(self, empty_blockchain: Blockchain, bt: BlockTools
     async def do_test_invalid_icc_sub_slot_vdf(
         self, keychain: Keychain, db_version: int, constants: ConsensusConstants
     ) -> None:
-        bt_high_iters = await create_block_tools_async(
-            constants=constants.replace(
-                SUB_SLOT_ITERS_STARTING=uint64(2**12),
-                DIFFICULTY_STARTING=uint64(2**14),
-            ),
-            keychain=keychain,
-        )
-        async with create_blockchain(bt_high_iters.constants, db_version) as (bc1, _):
+        async with (
+            create_block_tools_async(
+                constants=constants.replace(
+                    SUB_SLOT_ITERS_STARTING=uint64(2**12),
+                    DIFFICULTY_STARTING=uint64(2**14),
+                ),
+                keychain=keychain,
+            ) as bt_high_iters,
+            create_blockchain(bt_high_iters.constants, db_version) as (bc1, _),
+        ):
             blocks = bt_high_iters.get_consecutive_blocks(10)
             for block in blocks:
                 if (

chia/_tests/conftest.py

@@ -231,12 +231,12 @@ def blockchain_constants(consensus_mode: ConsensusMode) -> ConsensusConstants:
 
 
 @pytest.fixture(scope="session", name="bt")
-async def block_tools_fixture(get_keychain, blockchain_constants, anyio_backend, testrun_uid: str) -> BlockTools:
+async def block_tools_fixture(get_keychain, blockchain_constants, anyio_backend, testrun_uid: str):
     # Note that this causes a lot of CPU and disk traffic - disk, DB, ports, process creation ...
-    shared_block_tools = await create_block_tools_async(
+    async with create_block_tools_async(
         constants=blockchain_constants, keychain=get_keychain, testrun_uid=testrun_uid
-    )
-    return shared_block_tools
+    ) as shared_block_tools:
+        yield shared_block_tools
 
 
 # if you have a system that has an unusual hostname for localhost and you want
@@ -952,19 +952,20 @@ def get_temp_keyring():
 
 @pytest.fixture(scope="function")
 async def get_b_tools_1(get_temp_keyring, testrun_uid: str):
-    return await create_block_tools_async(
+    async with create_block_tools_async(
         constants=test_constants_modified, keychain=get_temp_keyring, testrun_uid=testrun_uid
-    )
+    ) as bt:
+        yield bt
 
 
 @pytest.fixture(scope="function")
 async def get_b_tools(get_temp_keyring, testrun_uid):
-    local_b_tools = await create_block_tools_async(
+    async with create_block_tools_async(
         constants=test_constants_modified, keychain=get_temp_keyring, testrun_uid=testrun_uid
-    )
-    new_config = local_b_tools._config
-    local_b_tools.change_config(new_config)
-    return local_b_tools
+    ) as local_b_tools:
+        new_config = local_b_tools._config
+        local_b_tools.change_config(new_config)
+        yield local_b_tools
 
 
 @pytest.fixture(scope="function")
@@ -1281,23 +1282,27 @@ async def farmer_harvester_2_simulators_zero_bits_plot_filter(
     )
 
     async with AsyncExitStack() as async_exit_stack:
-        bt = await create_block_tools_async(
-            zero_bit_plot_filter_consts,
-            keychain=get_temp_keyring,
-            testrun_uid=testrun_uid,
+        bt = await async_exit_stack.enter_async_context(
+            create_block_tools_async(
+                zero_bit_plot_filter_consts,
+                keychain=get_temp_keyring,
+                testrun_uid=testrun_uid,
+            )
         )
 
         config_overrides: dict[str, int] = {"full_node.max_sync_wait": 0}
 
         bts = [
-            await create_block_tools_async(
-                zero_bit_plot_filter_consts,
-                keychain=get_temp_keyring,
-                num_og_plots=0,
-                num_pool_plots=0,
-                num_non_keychain_plots=0,
-                config_overrides=config_overrides,
-                testrun_uid=testrun_uid,
+            await async_exit_stack.enter_async_context(
+                create_block_tools_async(
+                    zero_bit_plot_filter_consts,
+                    keychain=get_temp_keyring,
+                    num_og_plots=0,
+                    num_pool_plots=0,
+                    num_non_keychain_plots=0,
+                    config_overrides=config_overrides,
+                    testrun_uid=testrun_uid,
+                )
             )
             for _ in range(2)
         ]

chia/_tests/core/full_node/stores/test_full_node_store.py

@@ -38,7 +38,8 @@ async def custom_block_tools(blockchain_constants: ConsensusConstants) -> AsyncI
             DISCRIMINANT_SIZE_BITS=uint16(32),
             SUB_SLOT_ITERS_STARTING=uint64(2**12),
         )
-        yield await create_block_tools_async(constants=patched_constants, keychain=keychain)
+        async with create_block_tools_async(constants=patched_constants, keychain=keychain) as bt:
+            yield bt
 
 
 @pytest.fixture(scope="function")

chia/_tests/core/full_node/test_full_node.py

@@ -2448,31 +2448,30 @@ async def test_invalid_capability_can_connect(
 @pytest.mark.anyio
 async def test_node_start_with_existing_blocks(db_version: int) -> None:
     with TempKeyring(populate=True) as keychain:
-        block_tools = await create_block_tools_async(keychain=keychain)
-
-        blocks_per_cycle = 5
-        expected_height = 0
-
-        for cycle in range(2):
-            async with setup_full_node(
-                consensus_constants=block_tools.constants,
-                db_name="node_restart_test.db",
-                self_hostname=block_tools.config["self_hostname"],
-                local_bt=block_tools,
-                simulator=True,
-                db_version=db_version,
-                reuse_db=True,
-            ) as service:
-                simulator_api = service._api
-                assert isinstance(simulator_api, FullNodeSimulator)
-                await simulator_api.farm_blocks_to_puzzlehash(count=blocks_per_cycle)
-
-                expected_height += blocks_per_cycle
-                assert simulator_api.full_node._blockchain is not None
-                block_record = simulator_api.full_node._blockchain.get_peak()
-
-                assert block_record is not None, f"block_record is None on cycle {cycle + 1}"
-                assert block_record.height == expected_height, f"wrong height on cycle {cycle + 1}"
+        async with create_block_tools_async(keychain=keychain) as block_tools:
+            blocks_per_cycle = 5
+            expected_height = 0
+
+            for cycle in range(2):
+                async with setup_full_node(
+                    consensus_constants=block_tools.constants,
+                    db_name="node_restart_test.db",
+                    self_hostname=block_tools.config["self_hostname"],
+                    local_bt=block_tools,
+                    simulator=True,
+                    db_version=db_version,
+                    reuse_db=True,
+                ) as service:
+                    simulator_api = service._api
+                    assert isinstance(simulator_api, FullNodeSimulator)
+                    await simulator_api.farm_blocks_to_puzzlehash(count=blocks_per_cycle)
+
+                    expected_height += blocks_per_cycle
+                    assert simulator_api.full_node._blockchain is not None
+                    block_record = simulator_api.full_node._blockchain.get_peak()
+
+                    assert block_record is not None, f"block_record is None on cycle {cycle + 1}"
+                    assert block_record.height == expected_height, f"wrong height on cycle {cycle + 1}"
 
 
 @pytest.mark.anyio

chia/_tests/farmer_harvester/test_filter_prefix_bits.py

@@ -55,33 +55,33 @@ async def farmer_harvester_with_filter_size_9(
     async def have_connections() -> bool:
         return len(await farmer_rpc_cl.get_connections()) > 0
 
-    local_b_tools = await create_block_tools_async(
+    async with create_block_tools_async(
         constants=test_constants.replace(NUMBER_ZERO_BITS_PLOT_FILTER_V1=uint8(9)), keychain=get_temp_keyring
-    )
-    new_config = local_b_tools._config
-    local_b_tools.change_config(new_config)
-    async with setup_farmer_solver_multi_harvester(
-        local_b_tools, 1, tmp_path, local_b_tools.constants, start_services=True
-    ) as (harvesters, farmer_service, _):
-        harvester_service = harvesters[0]
-        assert farmer_service.rpc_server is not None
-        farmer_rpc_cl = await FarmerRpcClient.create(
-            self_hostname, farmer_service.rpc_server.listen_port, farmer_service.root_path, farmer_service.config
-        )
-        assert harvester_service.rpc_server is not None
-        harvester_rpc_cl = await HarvesterRpcClient.create(
-            self_hostname,
-            harvester_service.rpc_server.listen_port,
-            harvester_service.root_path,
-            harvester_service.config,
-        )
-        await time_out_assert(15, have_connections, True)
-        yield harvester_service, farmer_service._api
+    ) as local_b_tools:
+        new_config = local_b_tools._config
+        local_b_tools.change_config(new_config)
+        async with setup_farmer_solver_multi_harvester(
+            local_b_tools, 1, tmp_path, local_b_tools.constants, start_services=True
+        ) as (harvesters, farmer_service, _):
+            harvester_service = harvesters[0]
+            assert farmer_service.rpc_server is not None
+            farmer_rpc_cl = await FarmerRpcClient.create(
+                self_hostname, farmer_service.rpc_server.listen_port, farmer_service.root_path, farmer_service.config
+            )
+            assert harvester_service.rpc_server is not None
+            harvester_rpc_cl = await HarvesterRpcClient.create(
+                self_hostname,
+                harvester_service.rpc_server.listen_port,
+                harvester_service.root_path,
+                harvester_service.config,
+            )
+            await time_out_assert(15, have_connections, True)
+            yield harvester_service, farmer_service._api
 
-    farmer_rpc_cl.close()
-    harvester_rpc_cl.close()
-    await farmer_rpc_cl.await_closed()
-    await harvester_rpc_cl.await_closed()
+        farmer_rpc_cl.close()
+        harvester_rpc_cl.close()
+        await farmer_rpc_cl.await_closed()
+        await harvester_rpc_cl.await_closed()
 
 
 @pytest.mark.parametrize(argnames=["peak_height", "eligible_plots"], argvalues=[(5495999, 0), (5496000, 1)])

chia/_tests/simulation/test_simulation.py

@@ -58,15 +58,15 @@
 @pytest.fixture(scope="function")
 async def extra_node(self_hostname) -> AsyncIterator[FullNodeAPI | FullNodeSimulator]:
     with TempKeyring() as keychain:
-        b_tools = await create_block_tools_async(constants=test_constants_modified, keychain=keychain)
-        async with setup_full_node(
-            test_constants_modified,
-            "blockchain_test_3.db",
-            self_hostname,
-            b_tools,
-            db_version=2,
-        ) as service:
-            yield service._api
+        async with create_block_tools_async(constants=test_constants_modified, keychain=keychain) as b_tools:
+            async with setup_full_node(
+                test_constants_modified,
+                "blockchain_test_3.db",
+                self_hostname,
+                b_tools,
+                db_version=2,
+            ) as service:
+                yield service._api
 
 
 class FakeDNSResolver:

chia/_tests/solver/test_solver_service.py

@@ -18,8 +18,10 @@
 @pytest.mark.anyio
 async def test_solver_api_methods(blockchain_constants: ConsensusConstants, tmp_path: Path) -> None:
     with TempKeyring(populate=True) as keychain:
-        bt = await create_block_tools_async(constants=blockchain_constants, keychain=keychain)
-        async with setup_solver(tmp_path, bt, blockchain_constants) as solver_service:
+        async with (
+            create_block_tools_async(constants=blockchain_constants, keychain=keychain) as bt,
+            setup_solver(tmp_path, bt, blockchain_constants) as solver_service,
+        ):
             solver = solver_service._node
             solver_api = solver_service._api
             assert solver_api.ready() is True