Skip to content

Commit ec37e3a

Browse files
refactor: update test phase manager instance model
1 parent 7f1583d commit ec37e3a

File tree

4 files changed

+56
-49
lines changed

4 files changed

+56
-49
lines changed

src/ethereum_test_specs/blockchain.py

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -440,7 +440,7 @@ class BlockchainTest(BaseTest):
440440
genesis_environment: Environment = Field(default_factory=Environment)
441441
chain_id: int = 1
442442
exclude_full_post_state_in_output: bool = False
443-
test_phase_manager: TestPhaseManager | None = None
443+
test_phase_manager: TestPhaseManager | None = Field(default=None, exclude=True)
444444
"""
445445
Exclude the post state from the fixture output.
446446
In this case, the state verification is only performed based on the state root.
@@ -742,7 +742,6 @@ def _get_test_phase_blocks(self) -> List[Block]:
742742
for tx in self.test_phase_manager.execution_transactions:
743743
tx.test_phase = TestPhase.EXECUTION
744744
blocks.append(Block(txs=self.test_phase_manager.execution_transactions))
745-
746745
return blocks
747746

748747
def make_fixture(
@@ -760,9 +759,6 @@ def make_fixture(
760759
head = genesis.header.block_hash
761760
invalid_blocks = 0
762761

763-
if self.test_phase_manager is not None:
764-
self.blocks.extend(self._get_test_phase_blocks())
765-
766762
for i, block in enumerate(self.blocks):
767763
# This is the most common case, the RLP needs to be constructed
768764
# based on the transactions to be included in the block.
@@ -827,9 +823,7 @@ def make_hive_fixture(
827823
head_hash = genesis.header.block_hash
828824
invalid_blocks = 0
829825

830-
if self.test_phase_manager is not None:
831-
self.blocks.extend(self._get_test_phase_blocks())
832-
826+
built_block = None
833827
for i, block in enumerate(self.blocks):
834828
built_block = self.generate_block_data(
835829
t8n=t8n,
@@ -852,6 +846,10 @@ def make_hive_fixture(
852846
t8n, t8n_state=alloc, expected_state=block.expected_post_state
853847
)
854848
self.check_exception_test(exception=invalid_blocks > 0)
849+
850+
if built_block is None:
851+
raise Exception("No blocks to process in the test")
852+
855853
fcu_version = fork.engine_forkchoice_updated_version(
856854
built_block.header.number, built_block.header.timestamp
857855
)
@@ -931,6 +929,10 @@ def generate(
931929
) -> BaseFixture:
932930
"""Generate the BlockchainTest fixture."""
933931
t8n.reset_traces()
932+
933+
if self.test_phase_manager is not None:
934+
self.blocks.extend(self._get_test_phase_blocks())
935+
934936
if fixture_format in [
935937
BlockchainEngineFixture,
936938
BlockchainEngineXFixture,
Lines changed: 38 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,14 @@
11
"""Test phase management for Ethereum tests."""
22

33
from contextlib import contextmanager
4-
from contextvars import ContextVar
54
from enum import Enum
6-
from typing import Iterator, Optional
5+
from typing import Any, Iterator, List, Optional
76

8-
from pydantic import BaseModel
7+
from pydantic import GetCoreSchemaHandler
8+
from pydantic_core.core_schema import (
9+
PlainValidatorFunctionSchema,
10+
no_info_plain_validator_function,
11+
)
912

1013

1114
class TestPhase(Enum):
@@ -15,52 +18,44 @@ class TestPhase(Enum):
1518
EXECUTION = "execution"
1619

1720

18-
_current_phase: ContextVar[Optional[TestPhase]] = ContextVar("test_phase", default=TestPhase.SETUP)
19-
20-
21-
class TestPhaseManager(BaseModel):
21+
class TestPhaseManager:
2222
"""
2323
Manages test phases and collects transactions and blocks by phase.
2424
This class provides a mechanism for tracking "setup" and "execution" phases,
2525
Only supports "setup" and "execution" phases now.
2626
"""
2727

28-
model_config = {"arbitrary_types_allowed": True}
29-
30-
setup_transactions: list = []
31-
setup_blocks: list = []
32-
execution_transactions: list = []
33-
execution_blocks: list = []
34-
35-
def __init__(self, **data):
28+
def __init__(self, *args, **kwargs):
3629
"""Initialize the TestPhaseManager with empty transaction and block lists."""
37-
super().__init__(**data)
38-
self.setup_transactions = []
39-
self.setup_blocks = []
40-
self.execution_transactions = []
41-
self.execution_blocks = []
30+
self.setup_transactions: List = []
31+
self.setup_blocks: List = []
32+
self.execution_transactions: List = []
33+
self.execution_blocks: List = []
34+
self._current_phase: Optional[TestPhase] = TestPhase.EXECUTION
4235

4336
@contextmanager
4437
def setup(self) -> Iterator["TestPhaseManager"]:
4538
"""Context manager for the setup phase of a benchmark test."""
46-
token = _current_phase.set(TestPhase.SETUP)
39+
old_phase = self._current_phase
40+
self._current_phase = TestPhase.SETUP
4741
try:
4842
yield self
4943
finally:
50-
_current_phase.reset(token)
44+
self._current_phase = old_phase
5145

5246
@contextmanager
5347
def execution(self) -> Iterator["TestPhaseManager"]:
5448
"""Context manager for the execution phase of a test."""
55-
token = _current_phase.set(TestPhase.EXECUTION)
49+
old_phase = self._current_phase
50+
self._current_phase = TestPhase.EXECUTION
5651
try:
5752
yield self
5853
finally:
59-
_current_phase.reset(token)
54+
self._current_phase = old_phase
6055

6156
def add_transaction(self, tx) -> None:
6257
"""Add a transaction to the current phase."""
63-
current_phase = _current_phase.get()
58+
current_phase = self._current_phase
6459
tx.test_phase = current_phase
6560

6661
if current_phase == TestPhase.EXECUTION:
@@ -70,7 +65,7 @@ def add_transaction(self, tx) -> None:
7065

7166
def add_block(self, block) -> None:
7267
"""Add a block to the current phase."""
73-
current_phase = _current_phase.get()
68+
current_phase = self._current_phase
7469
for tx in block.txs:
7570
tx.test_phase = current_phase
7671

@@ -81,4 +76,20 @@ def add_block(self, block) -> None:
8176

8277
def get_current_phase(self) -> Optional[TestPhase]:
8378
"""Get the current test phase."""
84-
return _current_phase.get()
79+
return self._current_phase
80+
81+
@staticmethod
82+
def __get_pydantic_core_schema__(
83+
source_type: Any, handler: GetCoreSchemaHandler
84+
) -> PlainValidatorFunctionSchema:
85+
"""Pydantic schema for TestPhaseManager."""
86+
87+
def validate_test_phase_manager(value):
88+
"""Return the TestPhaseManager instance as-is."""
89+
if isinstance(value, source_type):
90+
return value
91+
return source_type()
92+
93+
return no_info_plain_validator_function(
94+
validate_test_phase_manager,
95+
)

src/pytest_plugins/shared/transaction_fixtures.py

Lines changed: 2 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -162,13 +162,7 @@ def typed_transaction(request, fork):
162162
) from e
163163

164164

165-
@pytest.fixture
166-
def phase_manager():
167-
"""Fixture providing access to a test phase manager."""
168-
return TestPhaseManager()
169-
170-
171-
@pytest.fixture
165+
@pytest.fixture(scope="function")
172166
def test_phase_manager():
173-
"""Fixture providing access to a test phase manager for blockchain tests."""
167+
"""Fixture providing access to a test phase manager."""
174168
return TestPhaseManager()

tests/benchmark/test_worst_stateful_opcodes.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -454,13 +454,13 @@ def test_worst_storage_access_warm(
454454
def test_worst_blockhash(
455455
blockchain_test: BlockchainTestFiller,
456456
pre: Alloc,
457-
phase_manager: TestPhaseManager,
457+
test_phase_manager: TestPhaseManager,
458458
gas_benchmark_value: int,
459459
):
460460
"""Test running a block with as many blockhash accessing oldest allowed block as possible."""
461-
with phase_manager.setup():
461+
with test_phase_manager.setup():
462462
for _ in range(256):
463-
phase_manager.add_block(Block())
463+
test_phase_manager.add_block(Block())
464464

465465
# Always ask for the oldest allowed BLOCKHASH block.
466466
execution_code = Op.PUSH1(1) + While(
@@ -473,13 +473,13 @@ def test_worst_blockhash(
473473
sender=pre.fund_eoa(),
474474
)
475475

476-
with phase_manager.execution():
477-
phase_manager.add_transaction(op_tx)
476+
with test_phase_manager.execution():
477+
test_phase_manager.add_transaction(op_tx)
478478

479479
blockchain_test(
480480
pre=pre,
481481
post={},
482-
test_phase_manager=phase_manager,
482+
test_phase_manager=test_phase_manager,
483483
)
484484

485485

0 commit comments

Comments
 (0)