refactor: remove benchmark manager feature

LouisTsai-Csie · LouisTsai-Csie · commit c489a91a486d · 2025-09-17T15:11:07.000+08:00
diff --git a/src/ethereum_test_specs/benchmark.py b/src/ethereum_test_specs/benchmark.py
@@ -1,19 +1,11 @@
 """Ethereum benchmark test spec definition and filler."""
 
 from abc import ABC, abstractmethod
-from contextlib import contextmanager
-from contextvars import ContextVar
 from dataclasses import dataclass, field
-from enum import Enum
-from typing import Any, Callable, ClassVar, Dict, Generator, List, Optional, Sequence, Type
+from typing import Callable, ClassVar, Dict, Generator, List, Sequence, Type
 
 import pytest
-from pydantic import ConfigDict, Field, GetCoreSchemaHandler
-from pydantic_core.core_schema import (
-    PlainValidatorFunctionSchema,
-    no_info_plain_validator_function,
-    to_string_ser_schema,
-)
+from pydantic import ConfigDict, Field
 
 from ethereum_clis import TransitionTool
 from ethereum_test_base_types import HexNumber
@@ -83,88 +75,6 @@ def _validate_code_size(self, code: Bytecode, fork: Fork) -> None:
             )
 
 
-class BenchmarkPhase(Enum):
-    """Phases of a benchmark test."""
-
-    SETUP = "setup"
-    EXECUTION = "execution"
-
-
-_current_phase: ContextVar[Optional[BenchmarkPhase]] = ContextVar("benchmark_phase", default=None)
-
-
-class BenchmarkManager:
-    """Context manager for managing benchmark test phases."""
-
-    def __init__(self):
-        """Initialize the BenchmarkManager with empty transaction and block lists."""
-        self.setup_transactions: List[Transaction] = []
-        self.setup_blocks: List[Block] = []
-        self.execution_transactions: List[Transaction] = []
-        self.execution_blocks: List[Block] = []
-
-    @contextmanager
-    def setup(self):
-        """Context manager for the setup phase of a benchmark test."""
-        token = _current_phase.set(BenchmarkPhase.SETUP)
-        try:
-            yield self
-        finally:
-            _current_phase.reset(token)
-
-    @contextmanager
-    def execution(self):
-        """Context manager for the execution phase of a benchmark test."""
-        token = _current_phase.set(BenchmarkPhase.EXECUTION)
-        try:
-            yield self
-        finally:
-            _current_phase.reset(token)
-
-    def add_transaction(self, tx: Transaction):
-        """Add a transaction to the current phase."""
-        current_phase = _current_phase.get()
-        if current_phase == BenchmarkPhase.SETUP:
-            self.setup_transactions.append(tx)
-        elif current_phase == BenchmarkPhase.EXECUTION:
-            self.execution_transactions.append(tx)
-        else:
-            self.setup_transactions.append(tx)
-
-    def add_block(self, block: Block):
-        """Add a block to the current phase."""
-        current_phase = _current_phase.get()
-        if current_phase == BenchmarkPhase.SETUP:
-            self.setup_blocks.append(block)
-        elif current_phase == BenchmarkPhase.EXECUTION:
-            self.execution_blocks.append(block)
-        else:
-            self.setup_blocks.append(block)
-
-    def get_current_phase(self) -> Optional[BenchmarkPhase]:
-        """Get the current benchmark phase."""
-        return _current_phase.get()
-
-    @classmethod
-    def __get_pydantic_core_schema__(
-        cls, source_type: Any, handler: GetCoreSchemaHandler
-    ) -> PlainValidatorFunctionSchema:
-        """Provide Pydantic core schema for BenchmarkManager serialization and validation."""
-
-        def validate_benchmark_manager(value):
-            if isinstance(value, cls):
-                return value
-            if value is None:
-                return None
-            # If value is passed as arguments, create new instance with no args
-            return cls()
-
-        return no_info_plain_validator_function(
-            validate_benchmark_manager,
-            serialization=to_string_ser_schema(),
-        )
-
-
 class BenchmarkTest(BaseTest):
     """Test type designed specifically for benchmark test cases."""
 
@@ -180,7 +90,6 @@ class BenchmarkTest(BaseTest):
     env: Environment = Field(default_factory=Environment)
     expected_benchmark_gas_used: int | None = None
     gas_benchmark_value: int = Field(default_factory=lambda: int(Environment().gas_limit))
-    benchmark_manager: BenchmarkManager | None = None
     code_generator: BenchmarkCodeGenerator | None = None
 
     supported_fixture_formats: ClassVar[Sequence[FixtureFormat | LabeledFixtureFormat]] = [
@@ -274,33 +183,6 @@ def generate_blockchain_test(self, fork: Fork) -> BlockchainTest:
                 blocks=generated_blocks,
             )
 
-        elif self.benchmark_manager is not None:
-            all_blocks = []
-            gas_limit = fork.transaction_gas_limit_cap() or self.gas_benchmark_value
-
-            if self.benchmark_manager.setup_blocks:
-                all_blocks.extend(self.benchmark_manager.setup_blocks)
-            elif self.benchmark_manager.setup_transactions:
-                setup_txs = []
-                for tx in self.benchmark_manager.setup_transactions:
-                    setup_txs.extend(self.split_transaction(tx, gas_limit))
-                all_blocks.append(Block(txs=setup_txs))
-
-            if self.benchmark_manager.execution_blocks:
-                all_blocks.extend(self.benchmark_manager.execution_blocks)
-            elif self.benchmark_manager.execution_transactions:
-                execution_txs = []
-                for tx in self.benchmark_manager.execution_transactions:
-                    execution_txs.extend(self.split_transaction(tx, gas_limit))
-                all_blocks.append(Block(txs=execution_txs))
-
-            return BlockchainTest.from_test(
-                base_test=self,
-                genesis_environment=self.env,
-                pre=self.pre,
-                post=self.post,
-                blocks=all_blocks,
-            )
         elif self.blocks is not None:
             return BlockchainTest.from_test(
                 base_test=self,
@@ -324,9 +206,7 @@ def generate_blockchain_test(self, fork: Fork) -> BlockchainTest:
                 genesis_environment=self.env,
             )
         else:
-            raise ValueError(
-                "Cannot create BlockchainTest without transactions, blocks, or benchmark_manager"
-            )
+            raise ValueError("Cannot create BlockchainTest without transactions or blocks")
 
     def generate(
         self,
@@ -358,10 +238,5 @@ def execute(
         raise Exception(f"Unsupported execute format: {execute_format}")
 
 
-def create_benchmark_manager() -> BenchmarkManager:
-    """Create a new BenchmarkManager instance for phase-aware benchmark testing."""
-    return BenchmarkManager()
-
-
 BenchmarkTestSpec = Callable[[str], Generator[BenchmarkTest, None, None]]
 BenchmarkTestFiller = Type[BenchmarkTest]
diff --git a/tests/benchmark/conftest.py b/tests/benchmark/conftest.py
@@ -5,7 +5,6 @@
 import pytest
 
 from ethereum_test_forks import Fork
-from ethereum_test_specs.benchmark import BenchmarkManager, create_benchmark_manager
 
 DEFAULT_BENCHMARK_FORK = "Prague"
 
@@ -68,9 +67,3 @@ def pytest_collection_modifyitems(config, items):
 def tx_gas_limit_cap(fork: Fork, gas_benchmark_value: int) -> int:
     """Return the transaction gas limit cap."""
     return fork.transaction_gas_limit_cap() or gas_benchmark_value
-
-
-@pytest.fixture
-def benchmark_manager() -> BenchmarkManager:
-    """Return a benchmark manager."""
-    return create_benchmark_manager()
diff --git a/tests/benchmark/test_worst_blocks.py b/tests/benchmark/test_worst_blocks.py
@@ -9,15 +9,14 @@
 
 import pytest
 
+from ethereum_test_base_types import Account
 from ethereum_test_forks import Fork
-from ethereum_test_specs.benchmark import BenchmarkManager
 from ethereum_test_tools import (
     AccessList,
-    Account,
     Address,
     Alloc,
-    BenchmarkTestFiller,
     Block,
+    BlockchainTestFiller,
     Environment,
     Hash,
     StateTestFiller,
@@ -112,8 +111,7 @@ def ether_transfer_case(
     ["a_to_a", "a_to_b", "diff_acc_to_b", "a_to_diff_acc", "diff_acc_to_diff_acc"],
 )
 def test_block_full_of_ether_transfers(
-    benchmark_test: BenchmarkTestFiller,
-    benchmark_manager: BenchmarkManager,
+    blockchain_test: BlockchainTestFiller,
     pre: Alloc,
     env: Environment,
     case_id: str,
@@ -138,18 +136,17 @@ def test_block_full_of_ether_transfers(
     # Create a single block with all transactions
     txs = []
     balances: dict[Address, int] = {}
-    with benchmark_manager.execution():
-        for _ in range(iteration_count):
-            receiver = next(receivers)
-            balances[receiver] = balances.get(receiver, 0) + transfer_amount
-            txs.append(
-                Transaction(
-                    to=receiver,
-                    value=transfer_amount,
-                    gas_limit=intrinsic_cost,
-                    sender=next(senders),
-                )
+    for _ in range(iteration_count):
+        receiver = next(receivers)
+        balances[receiver] = balances.get(receiver, 0) + transfer_amount
+        txs.append(
+            Transaction(
+                to=receiver,
+                value=transfer_amount,
+                gas_limit=intrinsic_cost,
+                sender=next(senders),
             )
+        )
 
     # Only include post state for non a_to_a cases
     post_state = (
@@ -158,8 +155,8 @@ def test_block_full_of_ether_transfers(
         else {receiver: Account(balance=balance) for receiver, balance in balances.items()}
     )
 
-    benchmark_test(
-        env=env,
+    blockchain_test(
+        genesis_environment=env,
         pre=pre,
         post=post_state,
         blocks=[Block(txs=txs)],
diff --git a/tests/benchmark/test_worst_stateful_opcodes.py b/tests/benchmark/test_worst_stateful_opcodes.py
@@ -10,8 +10,6 @@
 import pytest
 
 from ethereum_test_forks import Fork
-from ethereum_test_specs import BenchmarkTestFiller
-from ethereum_test_specs.benchmark import BenchmarkManager
 from ethereum_test_tools import (
     Account,
     Address,
@@ -453,35 +451,30 @@ def test_worst_storage_access_warm(
 
 
 def test_worst_blockhash(
-    benchmark_test: BenchmarkTestFiller,
-    benchmark_manager: BenchmarkManager,
+    blockchain_test: BlockchainTestFiller,
     pre: Alloc,
     gas_benchmark_value: int,
 ):
     """Test running a block with as many blockhash accessing oldest allowed block as possible."""
     # Create 256 dummy blocks to fill the blockhash window.
-    with benchmark_manager.setup():
-        for _ in range(256):
-            benchmark_manager.add_block(Block())
+    blocks = [Block()] * 256
 
     # Always ask for the oldest allowed BLOCKHASH block.
     execution_code = Op.PUSH1(1) + While(
         body=Op.POP(Op.BLOCKHASH(Op.DUP1)),
     )
     execution_code_address = pre.deploy_contract(code=execution_code)
-    with benchmark_manager.execution():
-        benchmark_manager.add_transaction(
-            Transaction(
-                to=execution_code_address,
-                gas_limit=gas_benchmark_value,
-                sender=pre.fund_eoa(),
-            )
-        )
+    op_tx = Transaction(
+        to=execution_code_address,
+        gas_limit=gas_benchmark_value,
+        sender=pre.fund_eoa(),
+    )
+    blocks.append(Block(txs=[op_tx]))
 
-    benchmark_test(
+    blockchain_test(
         pre=pre,
         post={},
-        benchmark_manager=benchmark_manager,
+        blocks=blocks,
     )
 
 
