1+ """Mocketizer decorator for managing Mocket lifecycle in tests."""
2+
3+ from __future__ import annotations
4+
5+ from typing import Any , Callable
6+
17from mocket .mocket import Mocket
28from mocket .mode import MocketMode
39from mocket .utils import get_mocketize
410
511
612class Mocketizer :
13+ """Context manager and decorator for managing Mocket lifecycle in tests."""
14+
715 def __init__ (
816 self ,
9- instance = None ,
10- namespace = None ,
11- truesocket_recording_dir = None ,
12- strict_mode = False ,
13- strict_mode_allowed = None ,
14- ):
17+ instance : Any | None = None ,
18+ namespace : str | None = None ,
19+ truesocket_recording_dir : str | None = None ,
20+ strict_mode : bool = False ,
21+ strict_mode_allowed : list | None = None ,
22+ ) -> None :
23+ """Initialize the Mocketizer.
24+
25+ Args:
26+ instance: Test instance (optional)
27+ namespace: Namespace for recordings
28+ truesocket_recording_dir: Directory for recording true socket calls
29+ strict_mode: Enable STRICT mode to forbid real socket calls
30+ strict_mode_allowed: List of allowed hosts in STRICT mode
31+ """
1532 self .instance = instance
1633 self .truesocket_recording_dir = truesocket_recording_dir
1734 self .namespace = namespace or str (id (self ))
@@ -23,41 +40,89 @@ def __init__(
2340 "Allowed locations are only accepted when STRICT mode is active."
2441 )
2542
26- def enter (self ):
43+ def enter (self ) -> None :
44+ """Enter the Mocketizer context (enable Mocket)."""
2745 Mocket .enable (
2846 namespace = self .namespace ,
2947 truesocket_recording_dir = self .truesocket_recording_dir ,
3048 )
3149 if self .instance :
3250 self .check_and_call ("mocketize_setup" )
3351
34- def __enter__ (self ):
52+ def __enter__ (self ) -> Mocketizer :
53+ """Enter context manager.
54+
55+ Returns:
56+ Self for use in `with` statements
57+ """
3558 self .enter ()
3659 return self
3760
38- def exit (self ):
61+ def exit (self ) -> None :
62+ """Exit the Mocketizer context (disable Mocket)."""
3963 if self .instance :
4064 self .check_and_call ("mocketize_teardown" )
4165
4266 Mocket .disable ()
4367
44- def __exit__ (self , type , value , tb ):
68+ def __exit__ (self , type : Any , value : Any , tb : Any ) -> None :
69+ """Exit context manager.
70+
71+ Args:
72+ type: Exception type
73+ value: Exception value
74+ tb: Traceback
75+ """
4576 self .exit ()
4677
47- async def __aenter__ (self , * args , ** kwargs ):
78+ async def __aenter__ (self , * args : Any , ** kwargs : Any ) -> Mocketizer :
79+ """Enter async context manager.
80+
81+ Returns:
82+ Self for use in `async with` statements
83+ """
4884 self .enter ()
4985 return self
5086
51- async def __aexit__ (self , * args , ** kwargs ):
87+ async def __aexit__ (self , * args : Any , ** kwargs : Any ) -> None :
88+ """Exit async context manager.
89+
90+ Args:
91+ *args: Exception arguments
92+ **kwargs: Exception keyword arguments
93+ """
5294 self .exit ()
5395
54- def check_and_call (self , method_name ):
96+ def check_and_call (self , method_name : str ) -> None :
97+ """Check if instance has a method and call it.
98+
99+ Args:
100+ method_name: Name of method to check and call
101+ """
55102 method = getattr (self .instance , method_name , None )
56103 if callable (method ):
57104 method ()
58105
59106 @staticmethod
60- def factory (test , truesocket_recording_dir , strict_mode , strict_mode_allowed , args ):
107+ def factory (
108+ test : Callable ,
109+ truesocket_recording_dir : str | None ,
110+ strict_mode : bool ,
111+ strict_mode_allowed : list | None ,
112+ args : tuple ,
113+ ) -> Mocketizer :
114+ """Create a Mocketizer instance for a test function.
115+
116+ Args:
117+ test: Test function being decorated
118+ truesocket_recording_dir: Recording directory
119+ strict_mode: Enable STRICT mode
120+ strict_mode_allowed: Allowed hosts in STRICT mode
121+ args: Positional arguments to test
122+
123+ Returns:
124+ Configured Mocketizer instance
125+ """
61126 instance = args [0 ] if args else None
62127 namespace = None
63128 if truesocket_recording_dir :
@@ -79,13 +144,26 @@ def factory(test, truesocket_recording_dir, strict_mode, strict_mode_allowed, ar
79144
80145
81146def wrapper (
82- test ,
83- truesocket_recording_dir = None ,
84- strict_mode = False ,
85- strict_mode_allowed = None ,
86- * args ,
87- ** kwargs ,
88- ):
147+ test : Callable ,
148+ truesocket_recording_dir : str | None = None ,
149+ strict_mode : bool = False ,
150+ strict_mode_allowed : list | None = None ,
151+ * args : Any ,
152+ ** kwargs : Any ,
153+ ) -> Any :
154+ """Wrapper function for @mocketize decorator.
155+
156+ Args:
157+ test: Test function to wrap
158+ truesocket_recording_dir: Recording directory
159+ strict_mode: Enable STRICT mode
160+ strict_mode_allowed: Allowed hosts in STRICT mode
161+ *args: Test arguments
162+ **kwargs: Test keyword arguments
163+
164+ Returns:
165+ Result of the test function
166+ """
89167 with Mocketizer .factory (
90168 test , truesocket_recording_dir , strict_mode , strict_mode_allowed , args
91169 ):
0 commit comments