11"""Test phase management for Ethereum tests.""" 
22
33from  contextlib  import  contextmanager 
4- from  contextvars  import  ContextVar 
54from  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
1114class  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+         )
0 commit comments