1
1
"""Test phase management for Ethereum tests."""
2
2
3
3
from contextlib import contextmanager
4
- from contextvars import ContextVar
5
4
from enum import Enum
6
- from typing import Iterator , Optional
5
+ from typing import Any , Iterator , List , Optional
7
6
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
+ )
9
12
10
13
11
14
class TestPhase (Enum ):
@@ -15,52 +18,44 @@ class TestPhase(Enum):
15
18
EXECUTION = "execution"
16
19
17
20
18
- _current_phase : ContextVar [Optional [TestPhase ]] = ContextVar ("test_phase" , default = TestPhase .SETUP )
19
-
20
-
21
- class TestPhaseManager (BaseModel ):
21
+ class TestPhaseManager :
22
22
"""
23
23
Manages test phases and collects transactions and blocks by phase.
24
24
This class provides a mechanism for tracking "setup" and "execution" phases,
25
25
Only supports "setup" and "execution" phases now.
26
26
"""
27
27
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 ):
36
29
"""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
42
35
43
36
@contextmanager
44
37
def setup (self ) -> Iterator ["TestPhaseManager" ]:
45
38
"""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
47
41
try :
48
42
yield self
49
43
finally :
50
- _current_phase . reset ( token )
44
+ self . _current_phase = old_phase
51
45
52
46
@contextmanager
53
47
def execution (self ) -> Iterator ["TestPhaseManager" ]:
54
48
"""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
56
51
try :
57
52
yield self
58
53
finally :
59
- _current_phase . reset ( token )
54
+ self . _current_phase = old_phase
60
55
61
56
def add_transaction (self , tx ) -> None :
62
57
"""Add a transaction to the current phase."""
63
- current_phase = _current_phase . get ()
58
+ current_phase = self . _current_phase
64
59
tx .test_phase = current_phase
65
60
66
61
if current_phase == TestPhase .EXECUTION :
@@ -70,7 +65,7 @@ def add_transaction(self, tx) -> None:
70
65
71
66
def add_block (self , block ) -> None :
72
67
"""Add a block to the current phase."""
73
- current_phase = _current_phase . get ()
68
+ current_phase = self . _current_phase
74
69
for tx in block .txs :
75
70
tx .test_phase = current_phase
76
71
@@ -81,4 +76,20 @@ def add_block(self, block) -> None:
81
76
82
77
def get_current_phase (self ) -> Optional [TestPhase ]:
83
78
"""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
+ )
0 commit comments