11"""Decorators for using the mock."""
22
33import re
4- import threading
54import time
65from contextlib import ContextDecorator
7- from typing import TYPE_CHECKING , Any , Literal , Self
6+ from typing import TYPE_CHECKING , Literal , Self
87from urllib .parse import urljoin , urlparse
98
10- import requests as requests_lib
9+ import requests
1110from beartype import BeartypeConf , beartype
1211from responses import RequestsMock
1312
2928 from collections .abc import Callable , Iterable , Mapping
3029
3130 from requests import PreparedRequest
32- from requests .adapters import HTTPAdapter # noqa: F401
3331
3432 ResponseType = tuple [int , Mapping [str , str ], str ]
35- Callback = Callable [[PreparedRequest ], ResponseType ] # noqa: F841
36-
37- # Thread-local storage to capture the request timeout
38- _timeout_storage = threading .local ()
33+ Callback = Callable [[PreparedRequest ], ResponseType ]
3934
4035_STRUCTURAL_SIMILARITY_MATCHER = StructuralSimilarityMatcher ()
4136_BRISQUE_TRACKING_RATER = BrisqueTargetTrackingRater ()
@@ -152,20 +147,28 @@ def __enter__(self) -> Self:
152147 def wrap_callback (callback : "Callback" ) -> "Callback" :
153148 """Wrap a callback to add a response delay."""
154149
155- def wrapped (request : "PreparedRequest" ) -> "ResponseType" :
156- # Check if the delay would exceed the request timeout
157- timeout = getattr (_timeout_storage , "timeout" , None )
158- if timeout is not None and delay_seconds > 0 :
159- # timeout can be a float or a tuple (connect, read)
160- if isinstance (timeout , tuple ):
161- effective_timeout : float | None = timeout [1 ] # read timeout
162- else :
163- effective_timeout = timeout
164- if (
165- effective_timeout is not None
166- and delay_seconds > effective_timeout
167- ):
168- raise requests_lib .exceptions .Timeout
150+ def wrapped (
151+ request : "PreparedRequest" ,
152+ ) -> "ResponseType" :
153+ # req_kwargs is added dynamically by the responses
154+ # library onto PreparedRequest objects - it is not
155+ # in the requests type stubs.
156+ timeout = request .req_kwargs .get ("timeout" ) # type: ignore[attr-defined]
157+ # requests allows timeout as a (connect, read)
158+ # tuple. The delay simulates server response
159+ # time, so compare against the read timeout.
160+ effective : float | None = None
161+ if isinstance (timeout , tuple ):
162+ effective = timeout [1 ]
163+ elif isinstance (timeout , (int , float )):
164+ effective = timeout
165+
166+ if (
167+ effective is not None
168+ and delay_seconds > effective
169+ ):
170+ time .sleep (effective )
171+ raise requests .exceptions .Timeout
169172
170173 result = callback (request )
171174 time .sleep (delay_seconds )
@@ -174,19 +177,6 @@ def wrapped(request: "PreparedRequest") -> "ResponseType":
174177 return wrapped
175178
176179 mock = RequestsMock (assert_all_requests_are_fired = False )
177-
178- # Patch _on_request to capture the timeout parameter
179- original_on_request = mock ._on_request # noqa: SLF001
180-
181- def patched_on_request (
182- adapter : "HTTPAdapter" ,
183- request : "PreparedRequest" ,
184- ** kwargs : Any , # noqa: ANN401
185- ) -> Any : # noqa: ANN401
186- _timeout_storage .timeout = kwargs .get ("timeout" )
187- return original_on_request (adapter , request , ** kwargs ) # type: ignore[misc]
188-
189- mock ._on_request = patched_on_request # type: ignore[method-assign] # noqa: SLF001
190180 for vws_route in self ._mock_vws_api .routes :
191181 url_pattern = urljoin (
192182 base = self ._base_vws_url ,
0 commit comments