41 lines
1.2 KiB
Python
41 lines
1.2 KiB
Python
|
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]]
|
||
|
|
|
||
|
|
|
||
|
|
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
|