492 lines
17 KiB
Python
492 lines
17 KiB
Python
import asyncio
|
||
import http
|
||
import inspect
|
||
import io
|
||
import json
|
||
import queue
|
||
import threading
|
||
import types
|
||
import typing
|
||
from urllib.parse import unquote, urljoin, urlsplit
|
||
|
||
import requests
|
||
|
||
from starlette.types import Message, Receive, Scope, Send
|
||
from starlette.websockets import WebSocketDisconnect
|
||
|
||
# Annotations for `Session.request()`
|
||
Cookies = typing.Union[
|
||
typing.MutableMapping[str, str], requests.cookies.RequestsCookieJar
|
||
]
|
||
Params = typing.Union[bytes, typing.MutableMapping[str, str]]
|
||
DataType = typing.Union[bytes, typing.MutableMapping[str, str], typing.IO]
|
||
TimeOut = typing.Union[float, typing.Tuple[float, float]]
|
||
FileType = typing.MutableMapping[str, typing.IO]
|
||
AuthType = typing.Union[
|
||
typing.Tuple[str, str],
|
||
requests.auth.AuthBase,
|
||
typing.Callable[[requests.Request], requests.Request],
|
||
]
|
||
|
||
|
||
ASGIInstance = typing.Callable[[Receive, Send], typing.Awaitable[None]]
|
||
ASGI2App = typing.Callable[[Scope], ASGIInstance]
|
||
ASGI3App = typing.Callable[[Scope, Receive, Send], typing.Awaitable[None]]
|
||
|
||
|
||
class _HeaderDict(requests.packages.urllib3._collections.HTTPHeaderDict):
|
||
def get_all(self, key: str, default: str) -> str:
|
||
return self.getheaders(key)
|
||
|
||
|
||
class _MockOriginalResponse:
|
||
"""
|
||
We have to jump through some hoops to present the response as if
|
||
it was made using urllib3.
|
||
"""
|
||
|
||
def __init__(self, headers: typing.List[typing.Tuple[bytes, bytes]]) -> None:
|
||
self.msg = _HeaderDict(headers)
|
||
self.closed = False
|
||
|
||
def isclosed(self) -> bool:
|
||
return self.closed
|
||
|
||
|
||
class _Upgrade(Exception):
|
||
def __init__(self, session: "WebSocketTestSession") -> None:
|
||
self.session = session
|
||
|
||
|
||
def _get_reason_phrase(status_code: int) -> str:
|
||
try:
|
||
return http.HTTPStatus(status_code).phrase
|
||
except ValueError:
|
||
return ""
|
||
|
||
|
||
def _is_asgi3(app: typing.Union[ASGI2App, ASGI3App]) -> bool:
|
||
if inspect.isclass(app):
|
||
return hasattr(app, "__await__")
|
||
elif inspect.isfunction(app):
|
||
return asyncio.iscoroutinefunction(app)
|
||
call = getattr(app, "__call__", None)
|
||
return asyncio.iscoroutinefunction(call)
|
||
|
||
|
||
class _WrapASGI2:
|
||
"""
|
||
Provide an ASGI3 interface onto an ASGI2 app.
|
||
"""
|
||
|
||
def __init__(self, app: ASGI2App) -> None:
|
||
self.app = app
|
||
|
||
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
|
||
instance = self.app(scope)
|
||
await instance(receive, send)
|
||
|
||
|
||
class _ASGIAdapter(requests.adapters.HTTPAdapter):
|
||
def __init__(
|
||
self, app: ASGI3App, raise_server_exceptions: bool = True, root_path: str = ""
|
||
) -> None:
|
||
self.app = app
|
||
self.raise_server_exceptions = raise_server_exceptions
|
||
self.root_path = root_path
|
||
|
||
def send(
|
||
self, request: requests.PreparedRequest, *args: typing.Any, **kwargs: typing.Any
|
||
) -> requests.Response:
|
||
scheme, netloc, path, query, fragment = (
|
||
str(item) for item in urlsplit(request.url)
|
||
)
|
||
|
||
default_port = {"http": 80, "ws": 80, "https": 443, "wss": 443}[scheme]
|
||
|
||
if ":" in netloc:
|
||
host, port_string = netloc.split(":", 1)
|
||
port = int(port_string)
|
||
else:
|
||
host = netloc
|
||
port = default_port
|
||
|
||
# Include the 'host' header.
|
||
if "host" in request.headers:
|
||
headers = [] # type: typing.List[typing.Tuple[bytes, bytes]]
|
||
elif port == default_port:
|
||
headers = [(b"host", host.encode())]
|
||
else:
|
||
headers = [(b"host", (f"{host}:{port}").encode())]
|
||
|
||
# Include other request headers.
|
||
headers += [
|
||
(key.lower().encode(), value.encode())
|
||
for key, value in request.headers.items()
|
||
]
|
||
|
||
if scheme in {"ws", "wss"}:
|
||
subprotocol = request.headers.get("sec-websocket-protocol", None)
|
||
if subprotocol is None:
|
||
subprotocols = [] # type: typing.Sequence[str]
|
||
else:
|
||
subprotocols = [value.strip() for value in subprotocol.split(",")]
|
||
scope = {
|
||
"type": "websocket",
|
||
"path": unquote(path),
|
||
"root_path": self.root_path,
|
||
"scheme": scheme,
|
||
"query_string": query.encode(),
|
||
"headers": headers,
|
||
"client": ["testclient", 50000],
|
||
"server": [host, port],
|
||
"subprotocols": subprotocols,
|
||
}
|
||
session = WebSocketTestSession(self.app, scope)
|
||
raise _Upgrade(session)
|
||
|
||
scope = {
|
||
"type": "http",
|
||
"http_version": "1.1",
|
||
"method": request.method,
|
||
"path": unquote(path),
|
||
"root_path": self.root_path,
|
||
"scheme": scheme,
|
||
"query_string": query.encode(),
|
||
"headers": headers,
|
||
"client": ["testclient", 50000],
|
||
"server": [host, port],
|
||
"extensions": {"http.response.template": {}},
|
||
}
|
||
|
||
request_complete = False
|
||
response_started = False
|
||
response_complete = False
|
||
raw_kwargs = {"body": io.BytesIO()} # type: typing.Dict[str, typing.Any]
|
||
template = None
|
||
context = None
|
||
|
||
async def receive() -> Message:
|
||
nonlocal request_complete, response_complete
|
||
|
||
if request_complete:
|
||
while not response_complete:
|
||
await asyncio.sleep(0.0001)
|
||
return {"type": "http.disconnect"}
|
||
|
||
body = request.body
|
||
if isinstance(body, str):
|
||
body_bytes = body.encode("utf-8") # type: bytes
|
||
elif body is None:
|
||
body_bytes = b""
|
||
elif isinstance(body, types.GeneratorType):
|
||
try:
|
||
chunk = body.send(None)
|
||
if isinstance(chunk, str):
|
||
chunk = chunk.encode("utf-8")
|
||
return {"type": "http.request", "body": chunk, "more_body": True}
|
||
except StopIteration:
|
||
request_complete = True
|
||
return {"type": "http.request", "body": b""}
|
||
else:
|
||
body_bytes = body
|
||
|
||
request_complete = True
|
||
return {"type": "http.request", "body": body_bytes}
|
||
|
||
async def send(message: Message) -> None:
|
||
nonlocal raw_kwargs, response_started, response_complete, template, context
|
||
|
||
if message["type"] == "http.response.start":
|
||
assert (
|
||
not response_started
|
||
), 'Received multiple "http.response.start" messages.'
|
||
raw_kwargs["version"] = 11
|
||
raw_kwargs["status"] = message["status"]
|
||
raw_kwargs["reason"] = _get_reason_phrase(message["status"])
|
||
raw_kwargs["headers"] = [
|
||
(key.decode(), value.decode()) for key, value in message["headers"]
|
||
]
|
||
raw_kwargs["preload_content"] = False
|
||
raw_kwargs["original_response"] = _MockOriginalResponse(
|
||
raw_kwargs["headers"]
|
||
)
|
||
response_started = True
|
||
elif message["type"] == "http.response.body":
|
||
assert (
|
||
response_started
|
||
), 'Received "http.response.body" without "http.response.start".'
|
||
assert (
|
||
not response_complete
|
||
), 'Received "http.response.body" after response completed.'
|
||
body = message.get("body", b"")
|
||
more_body = message.get("more_body", False)
|
||
if request.method != "HEAD":
|
||
raw_kwargs["body"].write(body)
|
||
if not more_body:
|
||
raw_kwargs["body"].seek(0)
|
||
response_complete = True
|
||
elif message["type"] == "http.response.template":
|
||
template = message["template"]
|
||
context = message["context"]
|
||
|
||
try:
|
||
loop = asyncio.get_event_loop()
|
||
except RuntimeError:
|
||
loop = asyncio.new_event_loop()
|
||
asyncio.set_event_loop(loop)
|
||
|
||
try:
|
||
loop.run_until_complete(self.app(scope, receive, send))
|
||
except BaseException as exc:
|
||
if self.raise_server_exceptions:
|
||
raise exc from None
|
||
|
||
if self.raise_server_exceptions:
|
||
assert response_started, "TestClient did not receive any response."
|
||
elif not response_started:
|
||
raw_kwargs = {
|
||
"version": 11,
|
||
"status": 500,
|
||
"reason": "Internal Server Error",
|
||
"headers": [],
|
||
"preload_content": False,
|
||
"original_response": _MockOriginalResponse([]),
|
||
"body": io.BytesIO(),
|
||
}
|
||
|
||
raw = requests.packages.urllib3.HTTPResponse(**raw_kwargs)
|
||
response = self.build_response(request, raw)
|
||
if template is not None:
|
||
response.template = template
|
||
response.context = context
|
||
return response
|
||
|
||
|
||
class WebSocketTestSession:
|
||
def __init__(self, app: ASGI3App, scope: Scope) -> None:
|
||
self.app = app
|
||
self.scope = scope
|
||
self.accepted_subprotocol = None
|
||
self._loop = asyncio.new_event_loop()
|
||
self._receive_queue = queue.Queue() # type: queue.Queue
|
||
self._send_queue = queue.Queue() # type: queue.Queue
|
||
self._thread = threading.Thread(target=self._run)
|
||
self.send({"type": "websocket.connect"})
|
||
self._thread.start()
|
||
message = self.receive()
|
||
self._raise_on_close(message)
|
||
self.accepted_subprotocol = message.get("subprotocol", None)
|
||
|
||
def __enter__(self) -> "WebSocketTestSession":
|
||
return self
|
||
|
||
def __exit__(self, *args: typing.Any) -> None:
|
||
self.close(1000)
|
||
self._thread.join()
|
||
while not self._send_queue.empty():
|
||
message = self._send_queue.get()
|
||
if isinstance(message, BaseException):
|
||
raise message
|
||
|
||
def _run(self) -> None:
|
||
"""
|
||
The sub-thread in which the websocket session runs.
|
||
"""
|
||
scope = self.scope
|
||
receive = self._asgi_receive
|
||
send = self._asgi_send
|
||
try:
|
||
self._loop.run_until_complete(self.app(scope, receive, send))
|
||
except BaseException as exc:
|
||
self._send_queue.put(exc)
|
||
|
||
async def _asgi_receive(self) -> Message:
|
||
while self._receive_queue.empty():
|
||
await asyncio.sleep(0)
|
||
return self._receive_queue.get()
|
||
|
||
async def _asgi_send(self, message: Message) -> None:
|
||
self._send_queue.put(message)
|
||
|
||
def _raise_on_close(self, message: Message) -> None:
|
||
if message["type"] == "websocket.close":
|
||
raise WebSocketDisconnect(message.get("code", 1000))
|
||
|
||
def send(self, message: Message) -> None:
|
||
self._receive_queue.put(message)
|
||
|
||
def send_text(self, data: str) -> None:
|
||
self.send({"type": "websocket.receive", "text": data})
|
||
|
||
def send_bytes(self, data: bytes) -> None:
|
||
self.send({"type": "websocket.receive", "bytes": data})
|
||
|
||
def send_json(self, data: typing.Any, mode: str = "text") -> None:
|
||
assert mode in ["text", "binary"]
|
||
text = json.dumps(data)
|
||
if mode == "text":
|
||
self.send({"type": "websocket.receive", "text": text})
|
||
else:
|
||
self.send({"type": "websocket.receive", "bytes": text.encode("utf-8")})
|
||
|
||
def close(self, code: int = 1000) -> None:
|
||
self.send({"type": "websocket.disconnect", "code": code})
|
||
|
||
def receive(self) -> Message:
|
||
message = self._send_queue.get()
|
||
if isinstance(message, BaseException):
|
||
raise message
|
||
return message
|
||
|
||
def receive_text(self) -> str:
|
||
message = self.receive()
|
||
self._raise_on_close(message)
|
||
return message["text"]
|
||
|
||
def receive_bytes(self) -> bytes:
|
||
message = self.receive()
|
||
self._raise_on_close(message)
|
||
return message["bytes"]
|
||
|
||
def receive_json(self, mode: str = "text") -> typing.Any:
|
||
assert mode in ["text", "binary"]
|
||
message = self.receive()
|
||
self._raise_on_close(message)
|
||
if mode == "text":
|
||
text = message["text"]
|
||
else:
|
||
text = message["bytes"].decode("utf-8")
|
||
return json.loads(text)
|
||
|
||
|
||
class TestClient(requests.Session):
|
||
__test__ = False # For pytest to not discover this up.
|
||
|
||
def __init__(
|
||
self,
|
||
app: typing.Union[ASGI2App, ASGI3App],
|
||
base_url: str = "http://testserver",
|
||
raise_server_exceptions: bool = True,
|
||
root_path: str = "",
|
||
) -> None:
|
||
super(TestClient, self).__init__()
|
||
if _is_asgi3(app):
|
||
app = typing.cast(ASGI3App, app)
|
||
asgi_app = app
|
||
else:
|
||
app = typing.cast(ASGI2App, app)
|
||
asgi_app = _WrapASGI2(app) # type: ignore
|
||
adapter = _ASGIAdapter(
|
||
asgi_app,
|
||
raise_server_exceptions=raise_server_exceptions,
|
||
root_path=root_path,
|
||
)
|
||
self.mount("http://", adapter)
|
||
self.mount("https://", adapter)
|
||
self.mount("ws://", adapter)
|
||
self.mount("wss://", adapter)
|
||
self.headers.update({"user-agent": "testclient"})
|
||
self.app = asgi_app
|
||
self.base_url = base_url
|
||
|
||
def request( # type: ignore
|
||
self,
|
||
method: str,
|
||
url: str,
|
||
params: Params = None,
|
||
data: DataType = None,
|
||
headers: typing.MutableMapping[str, str] = None,
|
||
cookies: Cookies = None,
|
||
files: FileType = None,
|
||
auth: AuthType = None,
|
||
timeout: TimeOut = None,
|
||
allow_redirects: bool = None,
|
||
proxies: typing.MutableMapping[str, str] = None,
|
||
hooks: typing.Any = None,
|
||
stream: bool = None,
|
||
verify: typing.Union[bool, str] = None,
|
||
cert: typing.Union[str, typing.Tuple[str, str]] = None,
|
||
json: typing.Any = None,
|
||
) -> requests.Response:
|
||
url = urljoin(self.base_url, url)
|
||
return super().request(
|
||
method,
|
||
url,
|
||
params=params,
|
||
data=data,
|
||
headers=headers,
|
||
cookies=cookies,
|
||
files=files,
|
||
auth=auth,
|
||
timeout=timeout,
|
||
allow_redirects=allow_redirects,
|
||
proxies=proxies,
|
||
hooks=hooks,
|
||
stream=stream,
|
||
verify=verify,
|
||
cert=cert,
|
||
json=json,
|
||
)
|
||
|
||
def websocket_connect(
|
||
self, url: str, subprotocols: typing.Sequence[str] = None, **kwargs: typing.Any
|
||
) -> typing.Any:
|
||
url = urljoin("ws://testserver", url)
|
||
headers = kwargs.get("headers", {})
|
||
headers.setdefault("connection", "upgrade")
|
||
headers.setdefault("sec-websocket-key", "testserver==")
|
||
headers.setdefault("sec-websocket-version", "13")
|
||
if subprotocols is not None:
|
||
headers.setdefault("sec-websocket-protocol", ", ".join(subprotocols))
|
||
kwargs["headers"] = headers
|
||
try:
|
||
super().request("GET", url, **kwargs)
|
||
except _Upgrade as exc:
|
||
session = exc.session
|
||
else:
|
||
raise RuntimeError("Expected WebSocket upgrade") # pragma: no cover
|
||
|
||
return session
|
||
|
||
def __enter__(self) -> requests.Session:
|
||
loop = asyncio.get_event_loop()
|
||
self.send_queue = asyncio.Queue() # type: asyncio.Queue
|
||
self.receive_queue = asyncio.Queue() # type: asyncio.Queue
|
||
self.task = loop.create_task(self.lifespan())
|
||
loop.run_until_complete(self.wait_startup())
|
||
return self
|
||
|
||
def __exit__(self, *args: typing.Any) -> None:
|
||
loop = asyncio.get_event_loop()
|
||
loop.run_until_complete(self.wait_shutdown())
|
||
|
||
async def lifespan(self) -> None:
|
||
scope = {"type": "lifespan"}
|
||
try:
|
||
await self.app(scope, self.receive_queue.get, self.send_queue.put)
|
||
finally:
|
||
await self.send_queue.put(None)
|
||
|
||
async def wait_startup(self) -> None:
|
||
await self.receive_queue.put({"type": "lifespan.startup"})
|
||
message = await self.send_queue.get()
|
||
if message is None:
|
||
self.task.result()
|
||
assert message["type"] in (
|
||
"lifespan.startup.complete",
|
||
"lifespan.startup.failed",
|
||
)
|
||
if message["type"] == "lifespan.startup.failed":
|
||
message = await self.send_queue.get()
|
||
if message is None:
|
||
self.task.result()
|
||
|
||
async def wait_shutdown(self) -> None:
|
||
await self.receive_queue.put({"type": "lifespan.shutdown"})
|
||
message = await self.send_queue.get()
|
||
if message is None:
|
||
self.task.result()
|
||
assert message["type"] == "lifespan.shutdown.complete"
|
||
await self.task
|