| 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 | |