Spaces:
Sleeping
Sleeping
from __future__ import annotations | |
import typing | |
from .._models import Request, Response | |
from .base import AsyncBaseTransport, BaseTransport | |
SyncHandler = typing.Callable[[Request], Response] | |
AsyncHandler = typing.Callable[[Request], typing.Coroutine[None, None, Response]] | |
__all__ = ["MockTransport"] | |
class MockTransport(AsyncBaseTransport, BaseTransport): | |
def __init__(self, handler: SyncHandler | AsyncHandler) -> None: | |
self.handler = handler | |
def handle_request( | |
self, | |
request: Request, | |
) -> Response: | |
request.read() | |
response = self.handler(request) | |
if not isinstance(response, Response): # pragma: no cover | |
raise TypeError("Cannot use an async handler in a sync Client") | |
return response | |
async def handle_async_request( | |
self, | |
request: Request, | |
) -> Response: | |
await request.aread() | |
response = self.handler(request) | |
# Allow handler to *optionally* be an `async` function. | |
# If it is, then the `response` variable need to be awaited to actually | |
# return the result. | |
if not isinstance(response, Response): | |
response = await response | |
return response | |