Skip to content

Commit 2e8ee6a

Browse files
refactor: update test phase manager instance model
1 parent 42a1e8c commit 2e8ee6a

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
@@ -423,7 +423,7 @@ class BlockchainTest(BaseTest):
423423
genesis_environment: Environment = Field(default_factory=Environment)
424424
chain_id: int = 1
425425
exclude_full_post_state_in_output: bool = False
426-
test_phase_manager: TestPhaseManager | None = None
426+
test_phase_manager: TestPhaseManager | None = Field(default=None, exclude=True)
427427
"""
428428
Exclude the post state from the fixture output. In this case, the state
429429
verification is only performed based on the state root.
@@ -736,7 +736,6 @@ def _get_test_phase_blocks(self) -> List[Block]:
736736
for tx in self.test_phase_manager.execution_transactions:
737737
tx.test_phase = TestPhase.EXECUTION
738738
blocks.append(Block(txs=self.test_phase_manager.execution_transactions))
739-
740739
return blocks
741740

742741
def make_fixture(
@@ -754,9 +753,6 @@ def make_fixture(
754753
head = genesis.header.block_hash
755754
invalid_blocks = 0
756755

757-
if self.test_phase_manager is not None:
758-
self.blocks.extend(self._get_test_phase_blocks())
759-
760756
for i, block in enumerate(self.blocks):
761757
# This is the most common case, the RLP needs to be constructed
762758
# based on the transactions to be included in the block.
@@ -822,9 +818,7 @@ def make_hive_fixture(
822818
head_hash = genesis.header.block_hash
823819
invalid_blocks = 0
824820

825-
if self.test_phase_manager is not None:
826-
self.blocks.extend(self._get_test_phase_blocks())
827-
821+
built_block = None
828822
for i, block in enumerate(self.blocks):
829823
built_block = self.generate_block_data(
830824
t8n=t8n,
@@ -847,6 +841,10 @@ def make_hive_fixture(
847841
t8n, t8n_state=alloc, expected_state=block.expected_post_state
848842
)
849843
self.check_exception_test(exception=invalid_blocks > 0)
844+
845+
if built_block is None:
846+
raise Exception("No blocks to process in the test")
847+
850848
fcu_version = fork.engine_forkchoice_updated_version(
851849
built_block.header.number, built_block.header.timestamp
852850
)
@@ -926,6 +924,10 @@ def generate(
926924
) -> BaseFixture:
927925
"""Generate the BlockchainTest fixture."""
928926
t8n.reset_traces()
927+
928+
if self.test_phase_manager is not None:
929+
self.blocks.extend(self._get_test_phase_blocks())
930+
929931
if fixture_format in [
930932
BlockchainEngineFixture,
931933
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
@@ -164,13 +164,7 @@ def typed_transaction(request, fork):
164164
) from e
165165

166166

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

tests/benchmark/test_worst_stateful_opcodes.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -463,16 +463,16 @@ def test_worst_storage_access_warm(
463463
def test_worst_blockhash(
464464
blockchain_test: BlockchainTestFiller,
465465
pre: Alloc,
466-
phase_manager: TestPhaseManager,
466+
test_phase_manager: TestPhaseManager,
467467
gas_benchmark_value: int,
468468
):
469469
"""
470470
Test running a block with as many blockhash accessing oldest allowed block
471471
as possible.
472472
"""
473-
with phase_manager.setup():
473+
with test_phase_manager.setup():
474474
for _ in range(256):
475-
phase_manager.add_block(Block())
475+
test_phase_manager.add_block(Block())
476476

477477
# Always ask for the oldest allowed BLOCKHASH block.
478478
execution_code = Op.PUSH1(1) + While(
@@ -485,13 +485,13 @@ def test_worst_blockhash(
485485
sender=pre.fund_eoa(),
486486
)
487487

488-
with phase_manager.execution():
489-
phase_manager.add_transaction(op_tx)
488+
with test_phase_manager.execution():
489+
test_phase_manager.add_transaction(op_tx)
490490

491491
blockchain_test(
492492
pre=pre,
493493
post={},
494-
test_phase_manager=phase_manager,
494+
test_phase_manager=test_phase_manager,
495495
)
496496

497497

0 commit comments

Comments
 (0)