Files
2026-10-21 13:04:35 +01:00

625 lines
24 KiB
Python

from __future__ import annotations
import contextlib
import inspect
import json
import math
import warnings
from collections.abc import Awaitable, Callable, Generator, Iterable, Mapping, Sequence
from concurrent.futures import Future, wait as wait_futures
from contextlib import AbstractContextManager
from types import GeneratorType
from typing import TYPE_CHECKING, Any, Literal, Self, TypedDict, TypeGuard, cast
from urllib.parse import unquote, urljoin
import anyio
import anyio.abc
import anyio.from_thread
from anyio.streams.stapled import StapledObjectStream
from starlette._utils import is_async_callable
from starlette.exceptions import StarletteDeprecationWarning
from starlette.types import ASGIApp, Message, Receive, Scope, Send
from starlette.websockets import WebSocketDisconnect
if TYPE_CHECKING:
import httpx2 as httpx
else:
try:
import httpx2 as httpx
except ModuleNotFoundError: # pragma: no cover
try:
import httpx
except ModuleNotFoundError:
raise RuntimeError(
"The starlette.testclient module requires the httpx2 package to be installed.\n"
"You can install this with:\n"
" $ pip install httpx2\n"
) from None
else:
warnings.warn(
"Using `httpx` with `starlette.testclient` is deprecated; install `httpx2` instead.",
StarletteDeprecationWarning,
stacklevel=2,
)
_PortalFactoryType = Callable[[], AbstractContextManager[anyio.from_thread.BlockingPortal]]
ASGIInstance = Callable[[Receive, Send], Awaitable[None]]
ASGI2App = Callable[[Scope], ASGIInstance]
ASGI3App = Callable[[Scope, Receive, Send], Awaitable[None]]
_RequestData = Mapping[str, str | Iterable[str] | bytes]
def _is_asgi3(app: ASGI2App | ASGI3App) -> TypeGuard[ASGI3App]:
if inspect.isclass(app):
return hasattr(app, "__await__")
return is_async_callable(app)
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 _AsyncBackend(TypedDict):
backend: str
backend_options: dict[str, Any]
class _Upgrade(Exception):
def __init__(self, session: WebSocketTestSession) -> None:
self.session = session
class WebSocketDenialResponse( # type: ignore[misc]
httpx.Response,
WebSocketDisconnect,
):
"""
A special case of `WebSocketDisconnect`, raised in the `TestClient` if the
`WebSocket` is closed before being accepted with a `send_denial_response()`.
"""
class WebSocketTestSession:
def __init__(
self,
app: ASGI3App,
scope: Scope,
portal_factory: _PortalFactoryType,
) -> None:
self.app = app
self.scope = scope
self.accepted_subprotocol = None
self.portal_factory = portal_factory
self.extra_headers = None
def __enter__(self) -> Self:
with contextlib.ExitStack() as stack:
self.portal = portal = stack.enter_context(self.portal_factory())
fut, cs = portal.start_task(self._run)
stack.callback(fut.result)
stack.callback(portal.call, cs.cancel)
self.send({"type": "websocket.connect"})
message = self.receive()
self._raise_on_close(message)
self.accepted_subprotocol = message.get("subprotocol", None)
self.extra_headers = message.get("headers", None)
stack.callback(self.close, 1000)
self.exit_stack = stack.pop_all()
return self
def __exit__(self, *args: Any) -> bool | None:
return self.exit_stack.__exit__(*args)
async def _run(self, *, task_status: anyio.abc.TaskStatus[anyio.CancelScope]) -> None:
"""
The sub-thread in which the websocket session runs.
"""
send_tx, send_rx = anyio.create_memory_object_stream[Message](math.inf)
receive_tx, receive_rx = anyio.create_memory_object_stream[Message](math.inf)
with send_tx, send_rx, receive_tx, receive_rx, anyio.CancelScope() as cs:
self._receive_tx = receive_tx
self._send_rx = send_rx
task_status.started(cs)
await self.app(self.scope, receive_rx.receive, send_tx.send)
# wait for cs.cancel to be called before closing streams
await anyio.sleep_forever()
def _raise_on_close(self, message: Message) -> None:
if message["type"] == "websocket.close":
raise WebSocketDisconnect(code=message.get("code", 1000), reason=message.get("reason", ""))
elif message["type"] == "websocket.http.response.start":
status_code: int = message["status"]
headers: list[tuple[bytes, bytes]] = message["headers"]
body: list[bytes] = []
while True:
message = self.receive()
assert message["type"] == "websocket.http.response.body"
body.append(message["body"])
if not message.get("more_body", False):
break
raise WebSocketDenialResponse(status_code=status_code, headers=headers, content=b"".join(body))
def send(self, message: Message) -> None:
self.portal.call(self._receive_tx.send, 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: Any, mode: Literal["text", "binary"] = "text") -> None:
text = json.dumps(data, separators=(",", ":"), ensure_ascii=False)
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, reason: str | None = None) -> None:
self.send({"type": "websocket.disconnect", "code": code, "reason": reason})
def receive(self) -> Message:
return self.portal.call(self._send_rx.receive)
def receive_text(self) -> str:
message = self.receive()
self._raise_on_close(message)
return cast(str, message["text"])
def receive_bytes(self) -> bytes:
message = self.receive()
self._raise_on_close(message)
return cast(bytes, message["bytes"])
def receive_json(self, mode: Literal["text", "binary"] = "text") -> Any:
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 _TestClientResponseStream(httpx.SyncByteStream):
def __init__(self, body: Generator[bytes, None, None], exit_stack: contextlib.ExitStack) -> None:
self._body = body
self._exit_stack = exit_stack
def __iter__(self) -> Generator[bytes, None, None]:
return self._body
def close(self) -> None:
self._exit_stack.close()
class _TestClientTransport(httpx.BaseTransport):
def __init__(
self,
app: ASGI3App,
portal_factory: _PortalFactoryType,
raise_server_exceptions: bool = True,
root_path: str = "",
*,
client: tuple[str, int],
app_state: dict[str, Any],
) -> None:
self.app = app
self.raise_server_exceptions = raise_server_exceptions
self.root_path = root_path
self.portal_factory = portal_factory
self.app_state = app_state
self.client = client
self._open_requests: set[Callable[[], None]] = set()
def close(self) -> None:
with contextlib.ExitStack() as stack:
while self._open_requests:
stack.callback(self._open_requests.pop())
def handle_request(self, request: httpx.Request) -> httpx.Response:
scheme = request.url.scheme
host = request.url.raw_host.decode(encoding="ascii")
path = request.url.path
raw_path = request.url.raw_path
query = request.url.query.decode(encoding="ascii")
default_port = {"http": 80, "ws": 80, "https": 443, "wss": 443}[scheme]
port = request.url.port
if port is None:
port = default_port
# Include the 'host' header.
if "host" in request.headers:
headers: list[tuple[bytes, bytes]] = []
else: # pragma: no cover
headers = [(b"host", request.url.netloc)]
# Include other request headers.
headers += [(key.lower().encode(), value.encode()) for key, value in request.headers.multi_items()]
scope: dict[str, Any]
if scheme in {"ws", "wss"}:
subprotocol = request.headers.get("sec-websocket-protocol", None)
if subprotocol is None:
subprotocols: Sequence[str] = []
else:
subprotocols = [value.strip() for value in subprotocol.split(",")]
scope = {
"type": "websocket",
"path": unquote(path),
"raw_path": raw_path.split(b"?", 1)[0],
"root_path": self.root_path,
"scheme": scheme,
"query_string": query.encode(),
"headers": headers,
"client": self.client,
"server": [host, port],
"subprotocols": subprotocols,
"state": self.app_state.copy(),
"extensions": {"websocket.http.response": {}},
}
session = WebSocketTestSession(self.app, scope, self.portal_factory)
raise _Upgrade(session)
scope = {
"type": "http",
"http_version": "1.1",
"method": request.method,
"path": unquote(path),
"raw_path": raw_path.split(b"?", 1)[0],
"root_path": self.root_path,
"scheme": scheme,
"query_string": query.encode(),
"headers": headers,
"client": self.client,
"server": [host, port],
"extensions": {"http.response.debug": {}, "http.response.trailers": {}},
"state": self.app_state.copy(),
}
request_complete = False
response_started = False
body_complete = False
trailers_expected = False
trailers: list[tuple[bytes, bytes]] = []
response_complete: anyio.Event
raw_kwargs: dict[str, Any] = {}
debug_info: dict[str, Any] | None = None
async def receive() -> Message:
nonlocal request_complete
if request_complete:
if not response_complete.is_set():
await response_complete.wait()
return {"type": "http.disconnect"}
body = request.read()
if isinstance(body, str):
body_bytes: bytes = body.encode("utf-8") # pragma: no cover
elif body is None:
body_bytes = b"" # pragma: no cover
elif isinstance(body, GeneratorType):
try: # pragma: no cover
chunk = body.send(None)
if isinstance(chunk, str):
chunk = chunk.encode("utf-8")
return {"type": "http.request", "body": chunk, "more_body": True}
except StopIteration: # pragma: no cover
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, debug_info, body_complete, trailers_expected
if message["type"] == "http.response.start":
assert not response_started, 'Received multiple "http.response.start" messages.'
raw_kwargs["status_code"] = message["status"]
raw_kwargs["headers"] = [(key.decode(), value.decode()) for key, value in message.get("headers", [])]
response_started = True
trailers_expected = message.get("trailers", False)
response_available.set()
elif message["type"] == "http.response.body":
assert response_started, 'Received "http.response.body" without "http.response.start".'
assert not response_complete.is_set(), 'Received "http.response.body" after response completed.'
assert not body_complete, 'Received "http.response.body" after body completed.'
body = message.get("body", b"")
more_body = message.get("more_body", False)
await body_tx.send((bytes(body), not more_body and not trailers_expected))
if not more_body:
body_complete = True
elif message["type"] == "http.response.trailers":
assert trailers_expected, 'Received "http.response.trailers" without declaring trailers.'
assert body_complete, 'Received "http.response.trailers" before body completed.'
assert not response_complete.is_set(), 'Received "http.response.trailers" after response completed.'
trailers.extend(message.get("headers", []))
if not message.get("more_trailers", False):
response_complete.set()
elif message["type"] == "http.response.debug":
debug_info = message["info"]
async def run_app() -> None:
try:
with cancel_scope:
async with body_tx:
await self.app(scope, receive, send)
finally:
response_available.set()
async def receive_body() -> bytes:
body, complete = await body_rx.receive()
if complete:
response_complete.set()
return body
def iter_body() -> Generator[bytes, None, None]:
while True:
try:
body = portal.call(receive_body)
except anyio.EndOfStream:
break
if request.method != "HEAD" and body:
yield body
exit_stack.close()
def cancel_if_incomplete() -> None:
if not response_complete.is_set():
cancel_scope.cancel()
def close_app(exc_type: type[BaseException] | None, *_: Any) -> None:
if app_task is None: # pragma: no cover - only if a signal interrupts start_task_soon()
portal.call(cancel_scope.cancel)
return
if not app_task.done():
portal.call(cancel_if_incomplete)
wait_futures([app_task])
self._open_requests.discard(abort)
if exc_type is None and self.raise_server_exceptions and (exception := app_task.exception()):
raise exception
def abort() -> None:
portal.call(cancel_scope.cancel)
exit_stack.close()
with contextlib.ExitStack() as stack:
portal = stack.enter_context(self.portal_factory())
response_available = portal.call(anyio.Event)
response_complete = portal.call(anyio.Event)
body_tx, body_rx = portal.call(anyio.create_memory_object_stream[tuple[bytes, bool]], 0)
stack.callback(body_tx.close)
stack.callback(body_rx.close)
cancel_scope = portal.call(anyio.CancelScope)
app_task: Future[None] | None = None
stack.push(close_app)
app_task = portal.start_task_soon(run_app)
portal.call(response_available.wait)
if not response_started:
stack.close()
if self.raise_server_exceptions:
raise AssertionError("TestClient did not receive any response.")
raw_kwargs = {"status_code": 500, "headers": [], "stream": httpx.ByteStream(b"")}
else:
exit_stack = stack.pop_all()
self._open_requests.add(abort)
raw_kwargs["stream"] = _TestClientResponseStream(iter_body(), exit_stack)
response = httpx.Response(**raw_kwargs, request=request)
if trailers_expected:
response.extensions["http.response.trailers"] = trailers
if debug_info is not None:
response.extensions["http.response.debug"] = debug_info
if "template" in debug_info:
response.template = debug_info["template"] # type: ignore[attr-defined]
if "context" in debug_info:
response.context = debug_info["context"] # type: ignore[attr-defined]
return response
class TestClient(httpx.Client):
__test__ = False
task: Future[None]
portal: anyio.from_thread.BlockingPortal | None = None
def __init__(
self,
app: ASGIApp,
base_url: str = "http://testserver",
raise_server_exceptions: bool = True,
root_path: str = "",
backend: Literal["asyncio", "trio"] = "asyncio",
backend_options: dict[str, Any] | None = None,
cookies: httpx._types.CookieTypes | None = None,
headers: dict[str, str] | None = None,
follow_redirects: bool = True,
client: tuple[str, int] = ("testclient", 50000),
) -> None:
self.async_backend = _AsyncBackend(backend=backend, backend_options=backend_options or {})
if _is_asgi3(app):
asgi_app = app
else:
app = cast(ASGI2App, app) # type: ignore[assignment]
asgi_app = _WrapASGI2(app) # type: ignore[arg-type]
self.app = asgi_app
self.app_state: dict[str, Any] = {}
transport = _TestClientTransport(
self.app,
portal_factory=self._portal_factory,
raise_server_exceptions=raise_server_exceptions,
root_path=root_path,
app_state=self.app_state,
client=client,
)
if headers is None:
headers = {}
headers.setdefault("user-agent", "testclient")
super().__init__(
base_url=base_url,
headers=headers,
transport=transport,
follow_redirects=follow_redirects,
cookies=cookies,
)
@contextlib.contextmanager
def _portal_factory(self) -> Generator[anyio.from_thread.BlockingPortal, None, None]:
if self.portal is not None:
yield self.portal
else:
with anyio.from_thread.start_blocking_portal(**self.async_backend) as portal:
yield portal
def request( # type: ignore[override]
self,
method: str,
url: httpx._types.URLTypes,
*,
content: httpx._types.RequestContent | None = None,
data: _RequestData | None = None,
files: httpx._types.RequestFiles | None = None,
json: Any = None,
params: httpx._types.QueryParamTypes | None = None,
headers: httpx._types.HeaderTypes | None = None,
cookies: httpx._types.CookieTypes | None = None,
auth: httpx._types.AuthTypes | httpx._client.UseClientDefault = httpx._client.USE_CLIENT_DEFAULT,
follow_redirects: bool | httpx._client.UseClientDefault = httpx._client.USE_CLIENT_DEFAULT,
timeout: httpx._types.TimeoutTypes | httpx._client.UseClientDefault = httpx._client.USE_CLIENT_DEFAULT,
extensions: dict[str, Any] | None = None,
) -> httpx.Response:
if timeout is not httpx.USE_CLIENT_DEFAULT:
warnings.warn(
"You should not use the 'timeout' argument with the TestClient. "
"See https://github.com/Kludex/starlette/issues/1108 for more information.",
StarletteDeprecationWarning,
stacklevel=2,
)
url = self._merge_url(url)
return super().request(
method,
url,
content=content,
data=data,
files=files,
json=json,
params=params,
headers=headers,
cookies=cookies,
auth=auth,
follow_redirects=follow_redirects,
timeout=timeout,
extensions=extensions,
)
def websocket_connect(
self,
url: str,
subprotocols: Sequence[str] | None = None,
**kwargs: Any,
) -> WebSocketTestSession:
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) -> Self:
with contextlib.ExitStack() as stack:
self.portal = portal = stack.enter_context(anyio.from_thread.start_blocking_portal(**self.async_backend))
@stack.callback
def reset_portal() -> None:
self.portal = None
send = anyio.create_memory_object_stream[Message | None](math.inf)
receive = anyio.create_memory_object_stream[Message](math.inf)
for channel in (*send, *receive):
stack.callback(channel.close)
self.stream_send = StapledObjectStream(*send)
self.stream_receive = StapledObjectStream(*receive)
self.task = portal.start_task_soon(self.lifespan)
portal.call(self.wait_startup)
@stack.callback
def wait_shutdown() -> None:
portal.call(self.wait_shutdown)
stack.callback(self._transport.close)
self.exit_stack = stack.pop_all()
return self
def __exit__(self, *args: Any) -> None:
self.exit_stack.close()
async def lifespan(self) -> None:
scope = {"type": "lifespan", "state": self.app_state}
try:
await self.app(scope, self.stream_receive.receive, self.stream_send.send)
finally:
await self.stream_send.send(None)
async def wait_startup(self) -> None:
await self.stream_receive.send({"type": "lifespan.startup"})
async def receive() -> Any:
message = await self.stream_send.receive()
if message is None:
self.task.result()
return message
message = await receive()
assert message["type"] in (
"lifespan.startup.complete",
"lifespan.startup.failed",
)
if message["type"] == "lifespan.startup.failed":
await receive()
async def wait_shutdown(self) -> None:
async def receive() -> Any:
message = await self.stream_send.receive()
if message is None:
self.task.result()
return message
await self.stream_receive.send({"type": "lifespan.shutdown"})
message = await receive()
assert message["type"] in (
"lifespan.shutdown.complete",
"lifespan.shutdown.failed",
)
if message["type"] == "lifespan.shutdown.failed":
await receive()