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