feat(tests): add phase manager to track testing phase

Author: LouisTsai-Csie

src/ethereum_test_execution/transaction_post.py

@@ -8,7 +8,7 @@
 from ethereum_test_base_types import Address, Alloc, Hash
 from ethereum_test_forks import Fork
 from ethereum_test_rpc import EngineRPC, EthRPC, SendTransactionExceptionError
-from ethereum_test_types import Transaction, TransactionTestMetadata
+from ethereum_test_types import TestPhase, Transaction, TransactionTestMetadata
 
 from .base import BaseExecute
 
@@ -40,9 +40,15 @@ def execute(
                 tx = tx.with_signature_and_sender()
                 to_address = tx.to
                 label = to_address.label if isinstance(to_address, Address) else None
+                phase = (
+                    "testing"
+                    if getattr(tx, "test_phase", None) == TestPhase.EXECUTION
+                    or getattr(block, "test_phase", None) == TestPhase.EXECUTION
+                    else "setup"
+                )
                 tx.metadata = TransactionTestMetadata(
                     test_id=request.node.nodeid,
-                    phase="testing",
+                    phase=phase,
                     target=label,
                     tx_index=tx_index,
                 )

src/ethereum_test_specs/blockchain.py

@@ -51,7 +51,16 @@
 )
 from ethereum_test_fixtures.common import FixtureBlobSchedule
 from ethereum_test_forks import Fork
-from ethereum_test_types import Alloc, Environment, Removable, Requests, Transaction, Withdrawal
+from ethereum_test_types import (
+    Alloc,
+    Environment,
+    Removable,
+    Requests,
+    TestPhase,
+    TestPhaseManager,
+    Transaction,
+    Withdrawal,
+)
 from ethereum_test_types.block_access_list import BlockAccessList, BlockAccessListExpectation
 
 from .base import BaseTest, OpMode, verify_result
@@ -232,7 +241,13 @@ class Block(Header):
     expected_post_state: Alloc | None = None
     """Post state for verification after block execution in BlockchainTest"""
     block_access_list: Bytes | None = Field(None)
-    """EIP-7928: Block-level access lists (serialized)."""
+    """
+        EIP-7928: Block-level access lists (serialized).
+    """
+    test_phase: TestPhase | None = None
+    """
+    Test phase for this block.
+    """
 
     def set_environment(self, env: Environment) -> Environment:
         """
@@ -404,10 +419,11 @@ class BlockchainTest(BaseTest):
 
     pre: Alloc
     post: Alloc
-    blocks: List[Block]
+    blocks: List[Block] = Field(default_factory=list)
     genesis_environment: Environment = Field(default_factory=Environment)
     chain_id: int = 1
     exclude_full_post_state_in_output: bool = False
+    test_phase_manager: TestPhaseManager | None = None
     """
     Exclude the post state from the fixture output. In this case, the state
     verification is only performed based on the state root.
@@ -698,6 +714,31 @@ def verify_post_state(self, t8n, t8n_state: Alloc, expected_state: Alloc | None
             print_traces(t8n.get_traces())
             raise e
 
+    def _get_test_phase_blocks(self) -> List[Block]:
+        """Get additional blocks from benchmark manager setup and execution phases."""
+        assert self.test_phase_manager is not None, "Test phase manager is not set"
+
+        blocks = []
+        if self.test_phase_manager.setup_blocks:
+            for block in self.test_phase_manager.setup_blocks:
+                block.test_phase = TestPhase.SETUP
+            blocks.extend(self.test_phase_manager.setup_blocks)
+        elif self.test_phase_manager.setup_transactions:
+            for tx in self.test_phase_manager.setup_transactions:
+                tx.test_phase = TestPhase.SETUP
+            blocks.append(Block(txs=self.test_phase_manager.setup_transactions))
+
+        if self.test_phase_manager.execution_blocks:
+            for block in self.test_phase_manager.execution_blocks:
+                block.test_phase = TestPhase.EXECUTION
+            blocks.extend(self.test_phase_manager.execution_blocks)
+        elif self.test_phase_manager.execution_transactions:
+            for tx in self.test_phase_manager.execution_transactions:
+                tx.test_phase = TestPhase.EXECUTION
+            blocks.append(Block(txs=self.test_phase_manager.execution_transactions))
+
+        return blocks
+
     def make_fixture(
         self,
         t8n: TransitionTool,
@@ -712,6 +753,10 @@ def make_fixture(
         env = environment_from_parent_header(genesis.header)
         head = genesis.header.block_hash
         invalid_blocks = 0
+
+        if self.test_phase_manager is not None:
+            self.blocks.extend(self._get_test_phase_blocks())
+
         for i, block in enumerate(self.blocks):
             # This is the most common case, the RLP needs to be constructed
             # based on the transactions to be included in the block.
@@ -776,6 +821,10 @@ def make_hive_fixture(
         env = environment_from_parent_header(genesis.header)
         head_hash = genesis.header.block_hash
         invalid_blocks = 0
+
+        if self.test_phase_manager is not None:
+            self.blocks.extend(self._get_test_phase_blocks())
+
         for i, block in enumerate(self.blocks):
             built_block = self.generate_block_data(
                 t8n=t8n,
@@ -897,6 +946,8 @@ def execute(
         """Generate the list of test fixtures."""
         if execute_format == TransactionPost:
             blocks: List[List[Transaction]] = []
+            if self.test_phase_manager is not None:
+                self.blocks.extend(self._get_test_phase_blocks())
             for block in self.blocks:
                 blocks += [block.txs]
             return TransactionPost(

src/ethereum_test_types/__init__.py

@@ -26,6 +26,7 @@
     compute_create_address,
     compute_eofcreate_address,
 )
+from .phase_manager import TestPhase, TestPhaseManager
 from .receipt_types import TransactionReceipt
 from .request_types import (
     ConsolidationRequest,
@@ -66,6 +67,8 @@
     "Removable",
     "Requests",
     "TestParameterGroup",
+    "TestPhase",
+    "TestPhaseManager",
     "Transaction",
     "TransactionDefaults",
     "TransactionReceipt",

src/ethereum_test_types/phase_manager.py

@@ -0,0 +1,84 @@
+"""Test phase management for Ethereum tests."""
+
+from contextlib import contextmanager
+from contextvars import ContextVar
+from enum import Enum
+from typing import Iterator, Optional
+
+from pydantic import BaseModel
+
+
+class TestPhase(Enum):
+    """Test phase for state and blockchain tests."""
+
+    SETUP = "setup"
+    EXECUTION = "execution"
+
+
+_current_phase: ContextVar[Optional[TestPhase]] = ContextVar("test_phase", default=TestPhase.SETUP)
+
+
+class TestPhaseManager(BaseModel):
+    """
+    Manages test phases and collects transactions and blocks by phase.
+    This class provides a mechanism for tracking "setup" and "execution" phases,
+    Only supports "setup" and "execution" phases now.
+    """
+
+    model_config = {"arbitrary_types_allowed": True}
+
+    setup_transactions: list = []
+    setup_blocks: list = []
+    execution_transactions: list = []
+    execution_blocks: list = []
+
+    def __init__(self, **data):
+        """Initialize the TestPhaseManager with empty transaction and block lists."""
+        super().__init__(**data)
+        self.setup_transactions = []
+        self.setup_blocks = []
+        self.execution_transactions = []
+        self.execution_blocks = []
+
+    @contextmanager
+    def setup(self) -> Iterator["TestPhaseManager"]:
+        """Context manager for the setup phase of a benchmark test."""
+        token = _current_phase.set(TestPhase.SETUP)
+        try:
+            yield self
+        finally:
+            _current_phase.reset(token)
+
+    @contextmanager
+    def execution(self) -> Iterator["TestPhaseManager"]:
+        """Context manager for the execution phase of a test."""
+        token = _current_phase.set(TestPhase.EXECUTION)
+        try:
+            yield self
+        finally:
+            _current_phase.reset(token)
+
+    def add_transaction(self, tx) -> None:
+        """Add a transaction to the current phase."""
+        current_phase = _current_phase.get()
+        tx.test_phase = current_phase
+
+        if current_phase == TestPhase.EXECUTION:
+            self.execution_transactions.append(tx)
+        else:
+            self.setup_transactions.append(tx)
+
+    def add_block(self, block) -> None:
+        """Add a block to the current phase."""
+        current_phase = _current_phase.get()
+        for tx in block.txs:
+            tx.test_phase = current_phase
+
+        if current_phase == TestPhase.EXECUTION:
+            self.execution_blocks.append(block)
+        else:
+            self.setup_blocks.append(block)
+
+    def get_current_phase(self) -> Optional[TestPhase]:
+        """Get the current test phase."""
+        return _current_phase.get()

src/ethereum_test_types/transaction_types.py

@@ -292,6 +292,7 @@ class Transaction(
     zero: ClassVar[Literal[0]] = 0
 
     metadata: TransactionTestMetadata | None = Field(None, exclude=True)
+    test_phase: str | None = None
 
     model_config = ConfigDict(validate_assignment=True)
 

src/pytest_plugins/shared/execute_fill.py

@@ -15,6 +15,7 @@
 from ..spec_version_checker.spec_version_checker import EIPSpecTestItem
 
 ALL_FIXTURE_PARAMETERS = {
+    "test_phase_manager",
     "gas_benchmark_value",
     "genesis_environment",
     "env",

src/pytest_plugins/shared/transaction_fixtures.py

@@ -9,6 +9,7 @@
 from ethereum_test_base_types import AccessList
 from ethereum_test_tools import Opcodes as Op
 from ethereum_test_types import AuthorizationTuple, Transaction, add_kzg_version
+from ethereum_test_types.phase_manager import TestPhaseManager
 
 
 @pytest.fixture
@@ -161,3 +162,15 @@ def typed_transaction(request, fork):
             f"Please add the missing fixture to "
             f"src/pytest_plugins/shared/transaction_fixtures.py"
         ) from e
+
+
+@pytest.fixture
+def phase_manager():
+    """Fixture providing access to a test phase manager."""
+    return TestPhaseManager()
+
+
+@pytest.fixture
+def test_phase_manager():
+    """Fixture providing access to a test phase manager for blockchain tests."""
+    return TestPhaseManager()

tests/benchmark/test_worst_stateful_opcodes.py

@@ -22,6 +22,7 @@
     compute_create2_address,
     compute_create_address,
 )
+from ethereum_test_types import TestPhaseManager
 from ethereum_test_vm import Opcodes as Op
 
 from .helpers import code_loop_precompile_call
@@ -462,14 +463,16 @@ def test_worst_storage_access_warm(
 def test_worst_blockhash(
     blockchain_test: BlockchainTestFiller,
     pre: Alloc,
+    phase_manager: TestPhaseManager,
     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.
-    blocks = [Block()] * 256
+    with phase_manager.setup():
+        for _ in range(256):
+            phase_manager.add_block(Block())
 
     # Always ask for the oldest allowed BLOCKHASH block.
     execution_code = Op.PUSH1(1) + While(
@@ -481,12 +484,14 @@ def test_worst_blockhash(
         gas_limit=gas_benchmark_value,
         sender=pre.fund_eoa(),
     )
-    blocks.append(Block(txs=[op_tx]))
+
+    with phase_manager.execution():
+        phase_manager.add_transaction(op_tx)
 
     blockchain_test(
         pre=pre,
         post={},
-        blocks=blocks,
+        test_phase_manager=phase_manager,
     )