switch from a benchmark mark to a fixture (#15190)

pytest.ini

@@ -8,7 +8,6 @@ log_format = %(asctime)s %(name)s: %(levelname)s %(message)s
 asyncio_mode = strict
 markers =
     limit_consensus_modes
-    benchmark
     data_layer: Mark as a data layer related test.
     test_mark_a1: used in testing test utilities
     test_mark_a2: used in testing test utilities

tests/conftest.py

@@ -61,6 +61,7 @@ from chia.wallet.wallet_node_api import WalletNodeAPI
 from tests.core.data_layer.util import ChiaRoot
 from tests.core.node_height import node_height_at_least
 from tests.simulation.test_simulation import test_constants_modified
+from tests.util.misc import BenchmarkRunner
 
 multiprocessing.set_start_method("spawn")
 
@@ -79,6 +80,12 @@ def seeded_random_fixture() -> random.Random:
     return seeded_random
 
 
+@pytest.fixture(name="benchmark_runner")
+def benchmark_runner_fixture(request: SubRequest) -> BenchmarkRunner:
+    label = request.node.name
+    return BenchmarkRunner(label=label)
+
+
 @pytest.fixture(name="node_name_for_file")
 def node_name_for_file_fixture(request: SubRequest) -> str:
     # TODO: handle other characters banned on windows
@@ -335,10 +342,15 @@ if os.getenv("_PYTEST_RAISE", "0") != "0":
         raise excinfo.value
 
 
+def pytest_configure(config):
+    config.addinivalue_line("markers", "benchmark: automatically assigned by the benchmark_runner fixture")
+
+
 def pytest_collection_modifyitems(session, config: pytest.Config, items: List[pytest.Function]):
     # https://github.com/pytest-dev/pytest/issues/3730#issuecomment-567142496
     removed = []
     kept = []
+    all_error_lines: List[str] = []
     limit_consensus_modes_problems: List[str] = []
     for item in items:
         limit_consensus_modes_marker = item.get_closest_marker("limit_consensus_modes")
@@ -364,8 +376,25 @@ def pytest_collection_modifyitems(session, config: pytest.Config, items: List[py
         items[:] = kept
 
     if len(limit_consensus_modes_problems) > 0:
-        name_lines = "\n".join(f"    {line}" for line in limit_consensus_modes_problems)
-        raise Exception(f"@pytest.mark.limit_consensus_modes used without consensus_mode:\n{name_lines}")
+        all_error_lines.append("@pytest.mark.limit_consensus_modes used without consensus_mode:")
+        all_error_lines.extend(f"    {line}" for line in limit_consensus_modes_problems)
+
+    benchmark_problems: List[str] = []
+    for item in items:
+        existing_benchmark_mark = item.get_closest_marker("benchmark")
+        if existing_benchmark_mark is not None:
+            benchmark_problems.append(item.name)
+
+        if "benchmark_runner" in getattr(item, "fixturenames", ()):
+            item.add_marker("benchmark")
+
+    if len(benchmark_problems) > 0:
+        all_error_lines.append("use the benchmark_runner fixture, not @pytest.mark.benchmark:")
+        all_error_lines.extend(f"    {line}" for line in benchmark_problems)
+
+    if len(all_error_lines) > 0:
+        all_error_lines.insert(0, "custom chia collection rules failed")
+        raise Exception("\n".join(all_error_lines))
 
 
 @pytest_asyncio.fixture(scope="function")

tests/core/data_layer/test_data_store.py

@@ -13,9 +13,6 @@ from typing import Any, Awaitable, Callable, Dict, List, Set, Tuple, cast
 import aiosqlite
 import pytest
 
-# TODO: update after resolution in https://github.com/pytest-dev/pytest/issues/7469
-from _pytest.fixtures import SubRequest
-
 from chia.data_layer.data_layer_errors import NodeHashError, TreeGenerationIncrementingError
 from chia.data_layer.data_layer_util import (
     DiffData,
@@ -46,7 +43,7 @@ from chia.types.blockchain_format.sized_bytes import bytes32
 from chia.util.byte_types import hexstr_to_bytes
 from chia.util.db_wrapper import DBWrapper2
 from tests.core.data_layer.util import Example, add_0123_example, add_01234567_example
-from tests.util.misc import Marks, assert_runtime, datacases
+from tests.util.misc import BenchmarkRunner, Marks, datacases
 
 log = logging.getLogger(__name__)
 
@@ -1402,12 +1399,11 @@ class BatchInsertBenchmarkCase:
         limit=24,
     ),
 )
-@pytest.mark.benchmark
 @pytest.mark.asyncio
 async def test_benchmark_batch_insert_speed(
     data_store: DataStore,
     tree_id: bytes32,
-    request: SubRequest,
+    benchmark_runner: BenchmarkRunner,
     case: BatchInsertBenchmarkCase,
 ) -> None:
     r = random.Random()
@@ -1432,7 +1428,7 @@ async def test_benchmark_batch_insert_speed(
             status=Status.COMMITTED,
         )
 
-    with assert_runtime(seconds=case.limit, label=request.node.name):
+    with benchmark_runner.assert_runtime(seconds=case.limit):
         await data_store.insert_batch(
             tree_id=tree_id,
             changelist=batch,

tests/core/full_node/test_node_load.py

@@ -6,12 +6,12 @@ from chia.simulator.time_out_assert import time_out_assert
 from chia.types.peer_info import PeerInfo
 from chia.util.ints import uint16
 from tests.connection_utils import connect_and_get_peer
-from tests.util.misc import assert_runtime
+from tests.util.misc import BenchmarkRunner
 
 
 class TestNodeLoad:
     @pytest.mark.asyncio
-    async def test_blocks_load(self, request: pytest.FixtureRequest, two_nodes, self_hostname):
+    async def test_blocks_load(self, two_nodes, self_hostname, benchmark_runner: BenchmarkRunner):
         num_blocks = 50
         full_node_1, full_node_2, server_1, server_2, bt = two_nodes
         blocks = bt.get_consecutive_blocks(num_blocks)
@@ -25,7 +25,7 @@ class TestNodeLoad:
 
         await time_out_assert(10, num_connections, 1)
 
-        with assert_runtime(seconds=100, label=request.node.name) as runtime_results_future:
+        with benchmark_runner.assert_runtime(seconds=100) as runtime_results_future:
             for i in range(1, num_blocks):
                 await full_node_1.full_node.add_block(blocks[i])
                 await full_node_2.full_node.add_block(blocks[i])

tests/core/full_node/test_performance.py

@@ -22,7 +22,7 @@ from chia.util.ints import uint64
 from tests.connection_utils import add_dummy_connection
 from tests.core.full_node.stores.test_coin_store import get_future_reward_coins
 from tests.core.node_height import node_height_at_least
-from tests.util.misc import assert_runtime
+from tests.util.misc import BenchmarkRunner
 
 log = logging.getLogger(__name__)
 
@@ -39,8 +39,9 @@ async def get_block_path(full_node: FullNodeAPI):
 
 class TestPerformance:
     @pytest.mark.asyncio
-    @pytest.mark.benchmark
-    async def test_full_block_performance(self, request: pytest.FixtureRequest, wallet_nodes_perf, self_hostname):
+    async def test_full_block_performance(
+        self, request: pytest.FixtureRequest, wallet_nodes_perf, self_hostname, benchmark_runner: BenchmarkRunner
+    ):
         full_node_1, server_1, wallet_a, wallet_receiver, bt = wallet_nodes_perf
         blocks = await full_node_1.get_all_full_blocks()
         full_node_1.full_node.mempool_manager.limit_factor = 1
@@ -118,7 +119,7 @@ class TestPerformance:
         pr = cProfile.Profile()
         pr.enable()
 
-        with assert_runtime(seconds=0.001, label=f"{request.node.name} - mempool"):
+        with benchmark_runner.assert_runtime(seconds=0.001, label=f"{request.node.name} - mempool"):
             num_tx: int = 0
             for spend_bundle, spend_bundle_id in zip(spend_bundles, spend_bundle_ids):
                 num_tx += 1
@@ -175,7 +176,7 @@ class TestPerformance:
         pr = cProfile.Profile()
         pr.enable()
 
-        with assert_runtime(seconds=0.1, label=f"{request.node.name} - unfinished"):
+        with benchmark_runner.assert_runtime(seconds=0.1, label=f"{request.node.name} - unfinished"):
             res = await full_node_1.respond_unfinished_block(fnp.RespondUnfinishedBlock(unfinished), fake_peer)
 
         log.warning(f"Res: {res}")
@@ -186,7 +187,7 @@ class TestPerformance:
         pr = cProfile.Profile()
         pr.enable()
 
-        with assert_runtime(seconds=0.1, label=f"{request.node.name} - full block"):
+        with benchmark_runner.assert_runtime(seconds=0.1, label=f"{request.node.name} - full block"):
             # No transactions generator, the full node already cached it from the unfinished block
             block_small = dataclasses.replace(block, transactions_generator=None)
             res = await full_node_1.full_node.add_block(block_small)

tests/core/mempool/test_mempool.py

@@ -59,7 +59,7 @@ from tests.core.mempool.test_mempool_manager import (
     spend_bundle_from_conditions,
 )
 from tests.core.node_height import node_height_at_least
-from tests.util.misc import assert_runtime
+from tests.util.misc import BenchmarkRunner
 
 BURN_PUZZLE_HASH = bytes32(b"0" * 32)
 BURN_PUZZLE_HASH_2 = bytes32(b"1" * 32)
@@ -2376,11 +2376,10 @@ class TestMaliciousGenerators:
             ConditionOpcode.ASSERT_SECONDS_RELATIVE,
         ],
     )
-    @pytest.mark.benchmark
-    def test_duplicate_large_integer_ladder(self, request, opcode, softfork_height):
+    def test_duplicate_large_integer_ladder(self, opcode, softfork_height, benchmark_runner: BenchmarkRunner):
         condition = SINGLE_ARG_INT_LADDER_COND.format(opcode=opcode.value[0], num=28, filler="0x00")
 
-        with assert_runtime(seconds=0.7, label=request.node.name):
+        with benchmark_runner.assert_runtime(seconds=0.7):
             npc_result = generator_condition_tester(condition, quote=False, height=softfork_height)
 
         assert npc_result.error == error_for_condition(opcode)
@@ -2394,11 +2393,10 @@ class TestMaliciousGenerators:
             ConditionOpcode.ASSERT_SECONDS_RELATIVE,
         ],
     )
-    @pytest.mark.benchmark
-    def test_duplicate_large_integer(self, request, opcode, softfork_height):
+    def test_duplicate_large_integer(self, opcode, softfork_height, benchmark_runner: BenchmarkRunner):
         condition = SINGLE_ARG_INT_COND.format(opcode=opcode.value[0], num=280000, val=100, filler="0x00")
 
-        with assert_runtime(seconds=1.1, label=request.node.name):
+        with benchmark_runner.assert_runtime(seconds=1.1):
             npc_result = generator_condition_tester(condition, quote=False, height=softfork_height)
 
         assert npc_result.error == error_for_condition(opcode)
@@ -2412,11 +2410,10 @@ class TestMaliciousGenerators:
             ConditionOpcode.ASSERT_SECONDS_RELATIVE,
         ],
     )
-    @pytest.mark.benchmark
-    def test_duplicate_large_integer_substr(self, request, opcode, softfork_height):
+    def test_duplicate_large_integer_substr(self, opcode, softfork_height, benchmark_runner: BenchmarkRunner):
         condition = SINGLE_ARG_INT_SUBSTR_COND.format(opcode=opcode.value[0], num=280000, val=100, filler="0x00")
 
-        with assert_runtime(seconds=1.5, label=request.node.name):
+        with benchmark_runner.assert_runtime(seconds=1.5):
             npc_result = generator_condition_tester(condition, quote=False, height=softfork_height)
 
         assert npc_result.error == error_for_condition(opcode)
@@ -2430,13 +2427,12 @@ class TestMaliciousGenerators:
             ConditionOpcode.ASSERT_SECONDS_RELATIVE,
         ],
     )
-    @pytest.mark.benchmark
-    def test_duplicate_large_integer_substr_tail(self, request, opcode, softfork_height):
+    def test_duplicate_large_integer_substr_tail(self, opcode, softfork_height, benchmark_runner: BenchmarkRunner):
         condition = SINGLE_ARG_INT_SUBSTR_TAIL_COND.format(
             opcode=opcode.value[0], num=280, val="0xffffffff", filler="0x00"
         )
 
-        with assert_runtime(seconds=0.3, label=request.node.name):
+        with benchmark_runner.assert_runtime(seconds=0.3):
             npc_result = generator_condition_tester(condition, quote=False, height=softfork_height)
 
         assert npc_result.error == error_for_condition(opcode)
@@ -2450,32 +2446,30 @@ class TestMaliciousGenerators:
             ConditionOpcode.ASSERT_SECONDS_RELATIVE,
         ],
     )
-    @pytest.mark.benchmark
-    def test_duplicate_large_integer_negative(self, request, opcode, softfork_height):
+    def test_duplicate_large_integer_negative(self, opcode, softfork_height, benchmark_runner: BenchmarkRunner):
         condition = SINGLE_ARG_INT_COND.format(opcode=opcode.value[0], num=280000, val=100, filler="0xff")
 
-        with assert_runtime(seconds=1, label=request.node.name):
+        with benchmark_runner.assert_runtime(seconds=1):
             npc_result = generator_condition_tester(condition, quote=False, height=softfork_height)
 
         assert npc_result.error is None
+        assert npc_result.conds is not None
         assert len(npc_result.conds.spends) == 1
 
-    @pytest.mark.benchmark
-    def test_duplicate_reserve_fee(self, request, softfork_height):
+    def test_duplicate_reserve_fee(self, softfork_height, benchmark_runner: BenchmarkRunner):
         opcode = ConditionOpcode.RESERVE_FEE
         condition = SINGLE_ARG_INT_COND.format(opcode=opcode.value[0], num=280000, val=100, filler="0x00")
 
-        with assert_runtime(seconds=1, label=request.node.name):
+        with benchmark_runner.assert_runtime(seconds=1):
             npc_result = generator_condition_tester(condition, quote=False, height=softfork_height)
 
         assert npc_result.error == error_for_condition(opcode)
 
-    @pytest.mark.benchmark
-    def test_duplicate_reserve_fee_negative(self, request: pytest.FixtureRequest, softfork_height):
+    def test_duplicate_reserve_fee_negative(self, softfork_height, benchmark_runner: BenchmarkRunner):
         opcode = ConditionOpcode.RESERVE_FEE
         condition = SINGLE_ARG_INT_COND.format(opcode=opcode.value[0], num=200000, val=100, filler="0xff")
 
-        with assert_runtime(seconds=0.8, label=request.node.name):
+        with benchmark_runner.assert_runtime(seconds=0.8):
             npc_result = generator_condition_tester(condition, quote=False, height=softfork_height)
 
         # RESERVE_FEE conditions fail unconditionally if they have a negative
@@ -2486,8 +2480,7 @@ class TestMaliciousGenerators:
     @pytest.mark.parametrize(
         "opcode", [ConditionOpcode.CREATE_COIN_ANNOUNCEMENT, ConditionOpcode.CREATE_PUZZLE_ANNOUNCEMENT]
     )
-    @pytest.mark.benchmark
-    def test_duplicate_coin_announces(self, request, opcode, softfork_height):
+    def test_duplicate_coin_announces(self, opcode, softfork_height, benchmark_runner: BenchmarkRunner):
         # with soft-fork3, we only allow 1024 create- or assert announcements
         # per spend
         if softfork_height >= test_constants.SOFT_FORK3_HEIGHT:
@@ -2495,42 +2488,42 @@ class TestMaliciousGenerators:
         else:
             condition = CREATE_ANNOUNCE_COND.format(opcode=opcode.value[0], num=5950000)
 
-        with assert_runtime(seconds=11, label=request.node.name):
+        with benchmark_runner.assert_runtime(seconds=11):
             npc_result = generator_condition_tester(condition, quote=False, height=softfork_height)
 
         assert npc_result.error is None
+        assert npc_result.conds is not None
         assert len(npc_result.conds.spends) == 1
         # coin announcements are not propagated to python, but validated in rust
         # TODO: optimize clvm to make this run in < 1 second
 
-    @pytest.mark.benchmark
-    def test_create_coin_duplicates(self, request: pytest.FixtureRequest, softfork_height):
+    def test_create_coin_duplicates(self, softfork_height, benchmark_runner: BenchmarkRunner):
         # CREATE_COIN
         # this program will emit 6000 identical CREATE_COIN conditions. However,
         # we'll just end up looking at two of them, and fail at the first
         # duplicate
         condition = CREATE_COIN.format(num=600000)
 
-        with assert_runtime(seconds=1.5, label=request.node.name):
+        with benchmark_runner.assert_runtime(seconds=1.5):
             npc_result = generator_condition_tester(condition, quote=False, height=softfork_height)
 
         assert npc_result.error == Err.DUPLICATE_OUTPUT.value
         assert npc_result.conds is None
 
-    @pytest.mark.benchmark
-    def test_many_create_coin(self, request, softfork_height):
+    def test_many_create_coin(self, softfork_height, benchmark_runner: BenchmarkRunner):
         # CREATE_COIN
         # this program will emit many CREATE_COIN conditions, all with different
         # amounts.
         # the number 6095 was chosen carefully to not exceed the maximum cost
         condition = CREATE_UNIQUE_COINS.format(num=6094)
 
-        with assert_runtime(seconds=0.3, label=request.node.name):
+        with benchmark_runner.assert_runtime(seconds=0.3):
             npc_result = generator_condition_tester(
                 condition, quote=False, height=softfork_height, coin_amount=123000000
             )
 
         assert npc_result.error is None
+        assert npc_result.conds is not None
         assert len(npc_result.conds.spends) == 1
         spend = npc_result.conds.spends[0]
         assert len(spend.create_coin) == 6094

tests/core/mempool/test_mempool_performance.py

@@ -14,7 +14,7 @@ from chia.wallet.transaction_record import TransactionRecord
 from chia.wallet.util.tx_config import DEFAULT_TX_CONFIG
 from chia.wallet.wallet_node import WalletNode
 from tests.connection_utils import connect_and_get_peer
-from tests.util.misc import assert_runtime
+from tests.util.misc import BenchmarkRunner
 
 
 async def wallet_height_at_least(wallet_node, h):
@@ -37,9 +37,8 @@ log = logging.getLogger(__name__)
 class TestMempoolPerformance:
     @pytest.mark.limit_consensus_modes(reason="benchmark")
     @pytest.mark.asyncio
-    @pytest.mark.benchmark
     async def test_mempool_update_performance(
-        self, request, wallet_nodes_mempool_perf, default_400_blocks, self_hostname
+        self, wallet_nodes_mempool_perf, default_400_blocks, self_hostname, benchmark_runner: BenchmarkRunner
     ):
         blocks = default_400_blocks
         full_nodes, wallets, bt = wallet_nodes_mempool_perf
@@ -83,5 +82,5 @@ class TestMempoolPerformance:
             else:
                 duration = 0.001
 
-            with assert_runtime(seconds=duration, label=request.node.name):
+            with benchmark_runner.assert_runtime(seconds=duration):
                 await full_node_api_1.full_node.add_block(block)

tests/core/test_cost_calculation.py

@@ -20,7 +20,7 @@ from chia.types.blockchain_format.serialized_program import SerializedProgram
 from chia.types.blockchain_format.sized_bytes import bytes32
 from chia.types.generator_types import BlockGenerator
 from chia.wallet.puzzles import p2_delegated_puzzle_or_hidden_puzzle
-from tests.util.misc import assert_runtime
+from tests.util.misc import BenchmarkRunner
 
 from .make_block_generator import make_block_generator
 
@@ -201,13 +201,12 @@ class TestCostCalculation:
         assert npc_result.error is None
 
     @pytest.mark.asyncio
-    @pytest.mark.benchmark
-    async def test_tx_generator_speed(self, request, softfork_height):
+    async def test_tx_generator_speed(self, softfork_height, benchmark_runner: BenchmarkRunner):
         LARGE_BLOCK_COIN_CONSUMED_COUNT = 687
         generator_bytes = large_block_generator(LARGE_BLOCK_COIN_CONSUMED_COUNT)
         program = SerializedProgram.from_bytes(generator_bytes)
 
-        with assert_runtime(seconds=0.5, label=request.node.name):
+        with benchmark_runner.assert_runtime(seconds=0.5):
             generator = BlockGenerator(program, [], [])
             npc_result = get_name_puzzle_conditions(
                 generator,
@@ -218,6 +217,7 @@ class TestCostCalculation:
             )
 
         assert npc_result.error is None
+        assert npc_result.conds is not None
         assert len(npc_result.conds.spends) == LARGE_BLOCK_COIN_CONSUMED_COUNT
 
     @pytest.mark.asyncio
@@ -252,8 +252,7 @@ class TestCostCalculation:
         assert npc_result.cost > 10000000
 
     @pytest.mark.asyncio
-    @pytest.mark.benchmark
-    async def test_standard_tx(self, request: pytest.FixtureRequest):
+    async def test_standard_tx(self, benchmark_runner: BenchmarkRunner):
         # this isn't a real public key, but we don't care
         public_key = bytes.fromhex(
             "af949b78fa6a957602c3593a3d6cb7711e08720415dad831ab18adacaa9b27ec3dda508ee32e24bc811c0abc5781ae21"
@@ -269,7 +268,7 @@ class TestCostCalculation:
             p2_delegated_puzzle_or_hidden_puzzle.solution_for_conditions(conditions)
         )
 
-        with assert_runtime(seconds=0.1, label=request.node.name):
+        with benchmark_runner.assert_runtime(seconds=0.1):
             total_cost = 0
             for i in range(0, 1000):
                 cost, result = puzzle_program.run_with_cost(test_constants.MAX_BLOCK_COST_CLVM, solution_program)
@@ -277,8 +276,7 @@ class TestCostCalculation:
 
 
 @pytest.mark.asyncio
-@pytest.mark.benchmark
-async def test_get_puzzle_and_solution_for_coin_performance():
+async def test_get_puzzle_and_solution_for_coin_performance(benchmark_runner: BenchmarkRunner):
     from clvm.casts import int_from_bytes
 
     from chia.full_node.mempool_check_conditions import DESERIALIZE_MOD
@@ -286,6 +284,7 @@ async def test_get_puzzle_and_solution_for_coin_performance():
 
     spends: List[Coin] = []
 
+    assert LARGE_BLOCK.transactions_generator is not None
     # first, list all spent coins in the block
     cost, result = LARGE_BLOCK.transactions_generator.run_with_cost(
         DEFAULT_CONSTANTS.MAX_BLOCK_COST_CLVM, DESERIALIZE_MOD, []
@@ -301,7 +300,7 @@ async def test_get_puzzle_and_solution_for_coin_performance():
     # benchmark the function to pick out the puzzle and solution for a specific
     # coin
     generator = BlockGenerator(LARGE_BLOCK.transactions_generator, [], [])
-    with assert_runtime(seconds=7, label="get_puzzle_and_solution_for_coin"):
+    with benchmark_runner.assert_runtime(seconds=7, label="get_puzzle_and_solution_for_coin"):
         for i in range(3):
             for c in spends:
                 spend_info = get_puzzle_and_solution_for_coin(generator, c, 0)

tests/util/misc.py

@@ -239,6 +239,7 @@ class _AssertRuntime:
     _results: Optional[AssertRuntimeResults] = None
     runtime_manager: Optional[contextlib.AbstractContextManager[Future[RuntimeResults]]] = None
     runtime_results_callable: Optional[Future[RuntimeResults]] = None
+    enable_assertion: bool = True
 
     def __enter__(self) -> Future[AssertRuntimeResults]:
         self.entry_line = caller_file_and_line()
@@ -281,15 +282,23 @@ class _AssertRuntime:
         if self.print:
             print(results.block(label=self.label))
 
-        if exc_type is None:
+        if exc_type is None and self.enable_assertion:
             __tracebackhide__ = True
             assert runtime.duration < self.seconds, results.message()
 
 
-# Related to the comment above about needing a class vs. using the context manager
-# decorator, this is just here to retain the function-style naming as the public
-# interface.  Hopefully we can switch away from the class at some point.
-assert_runtime = _AssertRuntime
+@final
+@dataclasses.dataclass
+class BenchmarkRunner:
+    enable_assertion: bool = True
+    label: Optional[str] = None
+
+    @functools.wraps(_AssertRuntime)
+    def assert_runtime(self, *args: Any, **kwargs: Any) -> _AssertRuntime:
+        kwargs.setdefault("enable_assertion", self.enable_assertion)
+        if self.label is not None:
+            kwargs.setdefault("label", self.label)
+        return _AssertRuntime(*args, **kwargs)
 
 
 @contextlib.contextmanager

tests/wallet/test_offer_parsing_performance.py

@@ -4,10 +4,8 @@ import cProfile
 from contextlib import contextmanager
 from typing import Iterator
 
-import pytest
-
 from chia.wallet.trading.offer import Offer
-from tests.util.misc import assert_runtime
+from tests.util.misc import BenchmarkRunner
 
 with_profile = False
 
@@ -28,21 +26,19 @@ def enable_profiler(name: str) -> Iterator[None]:
     pr.dump_stats(f"{name}.profile")
 
 
-@pytest.mark.benchmark
-def test_offer_parsing_performance() -> None:
+def test_offer_parsing_performance(benchmark_runner: BenchmarkRunner) -> None:
     offer_bytes = bytes.fromhex(test_offer)
-    with assert_runtime(seconds=2, label="Offer.from_bytes()"):
+    with benchmark_runner.assert_runtime(seconds=2, label="Offer.from_bytes()"):
         with enable_profiler("offer-parsing"):
             for _ in range(100):
                 o = Offer.from_bytes(offer_bytes)
                 assert o is not None
 
 
-@pytest.mark.benchmark
-def test_offered_coins_performance() -> None:
+def test_offered_coins_performance(benchmark_runner: BenchmarkRunner) -> None:
     offer_bytes = bytes.fromhex(test_offer)
     o = Offer.from_bytes(offer_bytes)
-    with assert_runtime(seconds=2.5, label="Offer.from_bytes()"):
+    with benchmark_runner.assert_runtime(seconds=2.5, label="Offer.from_bytes()"):
         with enable_profiler("offered-coins"):
             for _ in range(100):
                 c = o.get_offered_coins()