testclient.py 27 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739
  1. from __future__ import annotations
  2. import contextlib
  3. import inspect
  4. import io
  5. import json
  6. import math
  7. import sys
  8. import warnings
  9. from collections.abc import Awaitable, Callable, Generator, Iterable, Mapping, MutableMapping, Sequence
  10. from concurrent.futures import Future
  11. from contextlib import AbstractContextManager
  12. from types import GeneratorType
  13. from typing import (
  14. Any,
  15. Literal,
  16. TypedDict,
  17. TypeGuard,
  18. cast,
  19. )
  20. from urllib.parse import unquote, urljoin
  21. import anyio
  22. import anyio.abc
  23. import anyio.from_thread
  24. from anyio.streams.stapled import StapledObjectStream
  25. from starlette._utils import is_async_callable
  26. from starlette.types import ASGIApp, Message, Receive, Scope, Send
  27. from starlette.websockets import WebSocketDisconnect
  28. if sys.version_info >= (3, 11): # pragma: no cover
  29. from typing import Self
  30. else: # pragma: no cover
  31. from typing_extensions import Self
  32. try:
  33. import httpx
  34. except ModuleNotFoundError: # pragma: no cover
  35. raise RuntimeError(
  36. "The starlette.testclient module requires the httpx package to be installed.\n"
  37. "You can install this with:\n"
  38. " $ pip install httpx\n"
  39. )
  40. _PortalFactoryType = Callable[[], AbstractContextManager[anyio.abc.BlockingPortal]]
  41. ASGIInstance = Callable[[Receive, Send], Awaitable[None]]
  42. ASGI2App = Callable[[Scope], ASGIInstance]
  43. ASGI3App = Callable[[Scope, Receive, Send], Awaitable[None]]
  44. _RequestData = Mapping[str, str | Iterable[str] | bytes]
  45. def _is_asgi3(app: ASGI2App | ASGI3App) -> TypeGuard[ASGI3App]:
  46. if inspect.isclass(app):
  47. return hasattr(app, "__await__")
  48. return is_async_callable(app)
  49. class _WrapASGI2:
  50. """
  51. Provide an ASGI3 interface onto an ASGI2 app.
  52. """
  53. def __init__(self, app: ASGI2App) -> None:
  54. self.app = app
  55. async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
  56. instance = self.app(scope)
  57. await instance(receive, send)
  58. class _AsyncBackend(TypedDict):
  59. backend: str
  60. backend_options: dict[str, Any]
  61. class _Upgrade(Exception):
  62. def __init__(self, session: WebSocketTestSession) -> None:
  63. self.session = session
  64. class WebSocketDenialResponse( # type: ignore[misc]
  65. httpx.Response,
  66. WebSocketDisconnect,
  67. ):
  68. """
  69. A special case of `WebSocketDisconnect`, raised in the `TestClient` if the
  70. `WebSocket` is closed before being accepted with a `send_denial_response()`.
  71. """
  72. class WebSocketTestSession:
  73. def __init__(
  74. self,
  75. app: ASGI3App,
  76. scope: Scope,
  77. portal_factory: _PortalFactoryType,
  78. ) -> None:
  79. self.app = app
  80. self.scope = scope
  81. self.accepted_subprotocol = None
  82. self.portal_factory = portal_factory
  83. self.extra_headers = None
  84. def __enter__(self) -> WebSocketTestSession:
  85. with contextlib.ExitStack() as stack:
  86. self.portal = portal = stack.enter_context(self.portal_factory())
  87. fut, cs = portal.start_task(self._run)
  88. stack.callback(fut.result)
  89. stack.callback(portal.call, cs.cancel)
  90. self.send({"type": "websocket.connect"})
  91. message = self.receive()
  92. self._raise_on_close(message)
  93. self.accepted_subprotocol = message.get("subprotocol", None)
  94. self.extra_headers = message.get("headers", None)
  95. stack.callback(self.close, 1000)
  96. self.exit_stack = stack.pop_all()
  97. return self
  98. def __exit__(self, *args: Any) -> bool | None:
  99. return self.exit_stack.__exit__(*args)
  100. async def _run(self, *, task_status: anyio.abc.TaskStatus[anyio.CancelScope]) -> None:
  101. """
  102. The sub-thread in which the websocket session runs.
  103. """
  104. send: anyio.create_memory_object_stream[Message] = anyio.create_memory_object_stream(math.inf)
  105. send_tx, send_rx = send
  106. receive: anyio.create_memory_object_stream[Message] = anyio.create_memory_object_stream(math.inf)
  107. receive_tx, receive_rx = receive
  108. with send_tx, send_rx, receive_tx, receive_rx, anyio.CancelScope() as cs:
  109. self._receive_tx = receive_tx
  110. self._send_rx = send_rx
  111. task_status.started(cs)
  112. await self.app(self.scope, receive_rx.receive, send_tx.send)
  113. # wait for cs.cancel to be called before closing streams
  114. await anyio.sleep_forever()
  115. def _raise_on_close(self, message: Message) -> None:
  116. if message["type"] == "websocket.close":
  117. raise WebSocketDisconnect(code=message.get("code", 1000), reason=message.get("reason", ""))
  118. elif message["type"] == "websocket.http.response.start":
  119. status_code: int = message["status"]
  120. headers: list[tuple[bytes, bytes]] = message["headers"]
  121. body: list[bytes] = []
  122. while True:
  123. message = self.receive()
  124. assert message["type"] == "websocket.http.response.body"
  125. body.append(message["body"])
  126. if not message.get("more_body", False):
  127. break
  128. raise WebSocketDenialResponse(status_code=status_code, headers=headers, content=b"".join(body))
  129. def send(self, message: Message) -> None:
  130. self.portal.call(self._receive_tx.send, message)
  131. def send_text(self, data: str) -> None:
  132. self.send({"type": "websocket.receive", "text": data})
  133. def send_bytes(self, data: bytes) -> None:
  134. self.send({"type": "websocket.receive", "bytes": data})
  135. def send_json(self, data: Any, mode: Literal["text", "binary"] = "text") -> None:
  136. text = json.dumps(data, separators=(",", ":"), ensure_ascii=False)
  137. if mode == "text":
  138. self.send({"type": "websocket.receive", "text": text})
  139. else:
  140. self.send({"type": "websocket.receive", "bytes": text.encode("utf-8")})
  141. def close(self, code: int = 1000, reason: str | None = None) -> None:
  142. self.send({"type": "websocket.disconnect", "code": code, "reason": reason})
  143. def receive(self) -> Message:
  144. return self.portal.call(self._send_rx.receive)
  145. def receive_text(self) -> str:
  146. message = self.receive()
  147. self._raise_on_close(message)
  148. return cast(str, message["text"])
  149. def receive_bytes(self) -> bytes:
  150. message = self.receive()
  151. self._raise_on_close(message)
  152. return cast(bytes, message["bytes"])
  153. def receive_json(self, mode: Literal["text", "binary"] = "text") -> Any:
  154. message = self.receive()
  155. self._raise_on_close(message)
  156. if mode == "text":
  157. text = message["text"]
  158. else:
  159. text = message["bytes"].decode("utf-8")
  160. return json.loads(text)
  161. class _TestClientTransport(httpx.BaseTransport):
  162. def __init__(
  163. self,
  164. app: ASGI3App,
  165. portal_factory: _PortalFactoryType,
  166. raise_server_exceptions: bool = True,
  167. root_path: str = "",
  168. *,
  169. client: tuple[str, int],
  170. app_state: dict[str, Any],
  171. ) -> None:
  172. self.app = app
  173. self.raise_server_exceptions = raise_server_exceptions
  174. self.root_path = root_path
  175. self.portal_factory = portal_factory
  176. self.app_state = app_state
  177. self.client = client
  178. def handle_request(self, request: httpx.Request) -> httpx.Response:
  179. scheme = request.url.scheme
  180. netloc = request.url.netloc.decode(encoding="ascii")
  181. path = request.url.path
  182. raw_path = request.url.raw_path
  183. query = request.url.query.decode(encoding="ascii")
  184. default_port = {"http": 80, "ws": 80, "https": 443, "wss": 443}[scheme]
  185. if ":" in netloc:
  186. host, port_string = netloc.split(":", 1)
  187. port = int(port_string)
  188. else:
  189. host = netloc
  190. port = default_port
  191. # Include the 'host' header.
  192. if "host" in request.headers:
  193. headers: list[tuple[bytes, bytes]] = []
  194. elif port == default_port: # pragma: no cover
  195. headers = [(b"host", host.encode())]
  196. else: # pragma: no cover
  197. headers = [(b"host", (f"{host}:{port}").encode())]
  198. # Include other request headers.
  199. headers += [(key.lower().encode(), value.encode()) for key, value in request.headers.multi_items()]
  200. scope: dict[str, Any]
  201. if scheme in {"ws", "wss"}:
  202. subprotocol = request.headers.get("sec-websocket-protocol", None)
  203. if subprotocol is None:
  204. subprotocols: Sequence[str] = []
  205. else:
  206. subprotocols = [value.strip() for value in subprotocol.split(",")]
  207. scope = {
  208. "type": "websocket",
  209. "path": unquote(path),
  210. "raw_path": raw_path.split(b"?", 1)[0],
  211. "root_path": self.root_path,
  212. "scheme": scheme,
  213. "query_string": query.encode(),
  214. "headers": headers,
  215. "client": self.client,
  216. "server": [host, port],
  217. "subprotocols": subprotocols,
  218. "state": self.app_state.copy(),
  219. "extensions": {"websocket.http.response": {}},
  220. }
  221. session = WebSocketTestSession(self.app, scope, self.portal_factory)
  222. raise _Upgrade(session)
  223. scope = {
  224. "type": "http",
  225. "http_version": "1.1",
  226. "method": request.method,
  227. "path": unquote(path),
  228. "raw_path": raw_path.split(b"?", 1)[0],
  229. "root_path": self.root_path,
  230. "scheme": scheme,
  231. "query_string": query.encode(),
  232. "headers": headers,
  233. "client": self.client,
  234. "server": [host, port],
  235. "extensions": {"http.response.debug": {}},
  236. "state": self.app_state.copy(),
  237. }
  238. request_complete = False
  239. response_started = False
  240. response_complete: anyio.Event
  241. raw_kwargs: dict[str, Any] = {"stream": io.BytesIO()}
  242. template = None
  243. context = None
  244. async def receive() -> Message:
  245. nonlocal request_complete
  246. if request_complete:
  247. if not response_complete.is_set():
  248. await response_complete.wait()
  249. return {"type": "http.disconnect"}
  250. body = request.read()
  251. if isinstance(body, str):
  252. body_bytes: bytes = body.encode("utf-8") # pragma: no cover
  253. elif body is None:
  254. body_bytes = b"" # pragma: no cover
  255. elif isinstance(body, GeneratorType):
  256. try: # pragma: no cover
  257. chunk = body.send(None)
  258. if isinstance(chunk, str):
  259. chunk = chunk.encode("utf-8")
  260. return {"type": "http.request", "body": chunk, "more_body": True}
  261. except StopIteration: # pragma: no cover
  262. request_complete = True
  263. return {"type": "http.request", "body": b""}
  264. else:
  265. body_bytes = body
  266. request_complete = True
  267. return {"type": "http.request", "body": body_bytes}
  268. async def send(message: Message) -> None:
  269. nonlocal raw_kwargs, response_started, template, context
  270. if message["type"] == "http.response.start":
  271. assert not response_started, 'Received multiple "http.response.start" messages.'
  272. raw_kwargs["status_code"] = message["status"]
  273. raw_kwargs["headers"] = [(key.decode(), value.decode()) for key, value in message.get("headers", [])]
  274. response_started = True
  275. elif message["type"] == "http.response.body":
  276. assert response_started, 'Received "http.response.body" without "http.response.start".'
  277. assert not response_complete.is_set(), 'Received "http.response.body" after response completed.'
  278. body = message.get("body", b"")
  279. more_body = message.get("more_body", False)
  280. if request.method != "HEAD":
  281. raw_kwargs["stream"].write(body)
  282. if not more_body:
  283. raw_kwargs["stream"].seek(0)
  284. response_complete.set()
  285. elif message["type"] == "http.response.debug":
  286. template = message["info"]["template"]
  287. context = message["info"]["context"]
  288. try:
  289. with self.portal_factory() as portal:
  290. response_complete = portal.call(anyio.Event)
  291. portal.call(self.app, scope, receive, send)
  292. except BaseException as exc:
  293. if self.raise_server_exceptions:
  294. raise exc
  295. if self.raise_server_exceptions:
  296. assert response_started, "TestClient did not receive any response."
  297. elif not response_started:
  298. raw_kwargs = {
  299. "status_code": 500,
  300. "headers": [],
  301. "stream": io.BytesIO(),
  302. }
  303. raw_kwargs["stream"] = httpx.ByteStream(raw_kwargs["stream"].read())
  304. response = httpx.Response(**raw_kwargs, request=request)
  305. if template is not None:
  306. response.template = template # type: ignore[attr-defined]
  307. response.context = context # type: ignore[attr-defined]
  308. return response
  309. class TestClient(httpx.Client):
  310. __test__ = False
  311. task: Future[None]
  312. portal: anyio.abc.BlockingPortal | None = None
  313. def __init__(
  314. self,
  315. app: ASGIApp,
  316. base_url: str = "http://testserver",
  317. raise_server_exceptions: bool = True,
  318. root_path: str = "",
  319. backend: Literal["asyncio", "trio"] = "asyncio",
  320. backend_options: dict[str, Any] | None = None,
  321. cookies: httpx._types.CookieTypes | None = None,
  322. headers: dict[str, str] | None = None,
  323. follow_redirects: bool = True,
  324. client: tuple[str, int] = ("testclient", 50000),
  325. ) -> None:
  326. self.async_backend = _AsyncBackend(backend=backend, backend_options=backend_options or {})
  327. if _is_asgi3(app):
  328. asgi_app = app
  329. else:
  330. app = cast(ASGI2App, app) # type: ignore[assignment]
  331. asgi_app = _WrapASGI2(app) # type: ignore[arg-type]
  332. self.app = asgi_app
  333. self.app_state: dict[str, Any] = {}
  334. transport = _TestClientTransport(
  335. self.app,
  336. portal_factory=self._portal_factory,
  337. raise_server_exceptions=raise_server_exceptions,
  338. root_path=root_path,
  339. app_state=self.app_state,
  340. client=client,
  341. )
  342. if headers is None:
  343. headers = {}
  344. headers.setdefault("user-agent", "testclient")
  345. super().__init__(
  346. base_url=base_url,
  347. headers=headers,
  348. transport=transport,
  349. follow_redirects=follow_redirects,
  350. cookies=cookies,
  351. )
  352. @contextlib.contextmanager
  353. def _portal_factory(self) -> Generator[anyio.abc.BlockingPortal, None, None]:
  354. if self.portal is not None:
  355. yield self.portal
  356. else:
  357. with anyio.from_thread.start_blocking_portal(**self.async_backend) as portal:
  358. yield portal
  359. def request( # type: ignore[override]
  360. self,
  361. method: str,
  362. url: httpx._types.URLTypes,
  363. *,
  364. content: httpx._types.RequestContent | None = None,
  365. data: _RequestData | None = None,
  366. files: httpx._types.RequestFiles | None = None,
  367. json: Any = None,
  368. params: httpx._types.QueryParamTypes | None = None,
  369. headers: httpx._types.HeaderTypes | None = None,
  370. cookies: httpx._types.CookieTypes | None = None,
  371. auth: httpx._types.AuthTypes | httpx._client.UseClientDefault = httpx._client.USE_CLIENT_DEFAULT,
  372. follow_redirects: bool | httpx._client.UseClientDefault = httpx._client.USE_CLIENT_DEFAULT,
  373. timeout: httpx._types.TimeoutTypes | httpx._client.UseClientDefault = httpx._client.USE_CLIENT_DEFAULT,
  374. extensions: dict[str, Any] | None = None,
  375. ) -> httpx.Response:
  376. if timeout is not httpx.USE_CLIENT_DEFAULT:
  377. warnings.warn(
  378. "You should not use the 'timeout' argument with the TestClient. "
  379. "See https://github.com/Kludex/starlette/issues/1108 for more information.",
  380. DeprecationWarning,
  381. )
  382. url = self._merge_url(url)
  383. return super().request(
  384. method,
  385. url,
  386. content=content,
  387. data=data,
  388. files=files,
  389. json=json,
  390. params=params,
  391. headers=headers,
  392. cookies=cookies,
  393. auth=auth,
  394. follow_redirects=follow_redirects,
  395. timeout=timeout,
  396. extensions=extensions,
  397. )
  398. def get( # type: ignore[override]
  399. self,
  400. url: httpx._types.URLTypes,
  401. *,
  402. params: httpx._types.QueryParamTypes | None = None,
  403. headers: httpx._types.HeaderTypes | None = None,
  404. cookies: httpx._types.CookieTypes | None = None,
  405. auth: httpx._types.AuthTypes | httpx._client.UseClientDefault = httpx._client.USE_CLIENT_DEFAULT,
  406. follow_redirects: bool | httpx._client.UseClientDefault = httpx._client.USE_CLIENT_DEFAULT,
  407. timeout: httpx._types.TimeoutTypes | httpx._client.UseClientDefault = httpx._client.USE_CLIENT_DEFAULT,
  408. extensions: dict[str, Any] | None = None,
  409. ) -> httpx.Response:
  410. return super().get(
  411. url,
  412. params=params,
  413. headers=headers,
  414. cookies=cookies,
  415. auth=auth,
  416. follow_redirects=follow_redirects,
  417. timeout=timeout,
  418. extensions=extensions,
  419. )
  420. def options( # type: ignore[override]
  421. self,
  422. url: httpx._types.URLTypes,
  423. *,
  424. params: httpx._types.QueryParamTypes | None = None,
  425. headers: httpx._types.HeaderTypes | None = None,
  426. cookies: httpx._types.CookieTypes | None = None,
  427. auth: httpx._types.AuthTypes | httpx._client.UseClientDefault = httpx._client.USE_CLIENT_DEFAULT,
  428. follow_redirects: bool | httpx._client.UseClientDefault = httpx._client.USE_CLIENT_DEFAULT,
  429. timeout: httpx._types.TimeoutTypes | httpx._client.UseClientDefault = httpx._client.USE_CLIENT_DEFAULT,
  430. extensions: dict[str, Any] | None = None,
  431. ) -> httpx.Response:
  432. return super().options(
  433. url,
  434. params=params,
  435. headers=headers,
  436. cookies=cookies,
  437. auth=auth,
  438. follow_redirects=follow_redirects,
  439. timeout=timeout,
  440. extensions=extensions,
  441. )
  442. def head( # type: ignore[override]
  443. self,
  444. url: httpx._types.URLTypes,
  445. *,
  446. params: httpx._types.QueryParamTypes | None = None,
  447. headers: httpx._types.HeaderTypes | None = None,
  448. cookies: httpx._types.CookieTypes | None = None,
  449. auth: httpx._types.AuthTypes | httpx._client.UseClientDefault = httpx._client.USE_CLIENT_DEFAULT,
  450. follow_redirects: bool | httpx._client.UseClientDefault = httpx._client.USE_CLIENT_DEFAULT,
  451. timeout: httpx._types.TimeoutTypes | httpx._client.UseClientDefault = httpx._client.USE_CLIENT_DEFAULT,
  452. extensions: dict[str, Any] | None = None,
  453. ) -> httpx.Response:
  454. return super().head(
  455. url,
  456. params=params,
  457. headers=headers,
  458. cookies=cookies,
  459. auth=auth,
  460. follow_redirects=follow_redirects,
  461. timeout=timeout,
  462. extensions=extensions,
  463. )
  464. def post( # type: ignore[override]
  465. self,
  466. url: httpx._types.URLTypes,
  467. *,
  468. content: httpx._types.RequestContent | None = None,
  469. data: _RequestData | None = None,
  470. files: httpx._types.RequestFiles | None = None,
  471. json: Any = None,
  472. params: httpx._types.QueryParamTypes | None = None,
  473. headers: httpx._types.HeaderTypes | None = None,
  474. cookies: httpx._types.CookieTypes | None = None,
  475. auth: httpx._types.AuthTypes | httpx._client.UseClientDefault = httpx._client.USE_CLIENT_DEFAULT,
  476. follow_redirects: bool | httpx._client.UseClientDefault = httpx._client.USE_CLIENT_DEFAULT,
  477. timeout: httpx._types.TimeoutTypes | httpx._client.UseClientDefault = httpx._client.USE_CLIENT_DEFAULT,
  478. extensions: dict[str, Any] | None = None,
  479. ) -> httpx.Response:
  480. return super().post(
  481. url,
  482. content=content,
  483. data=data,
  484. files=files,
  485. json=json,
  486. params=params,
  487. headers=headers,
  488. cookies=cookies,
  489. auth=auth,
  490. follow_redirects=follow_redirects,
  491. timeout=timeout,
  492. extensions=extensions,
  493. )
  494. def put( # type: ignore[override]
  495. self,
  496. url: httpx._types.URLTypes,
  497. *,
  498. content: httpx._types.RequestContent | None = None,
  499. data: _RequestData | None = None,
  500. files: httpx._types.RequestFiles | None = None,
  501. json: Any = None,
  502. params: httpx._types.QueryParamTypes | None = None,
  503. headers: httpx._types.HeaderTypes | None = None,
  504. cookies: httpx._types.CookieTypes | None = None,
  505. auth: httpx._types.AuthTypes | httpx._client.UseClientDefault = httpx._client.USE_CLIENT_DEFAULT,
  506. follow_redirects: bool | httpx._client.UseClientDefault = httpx._client.USE_CLIENT_DEFAULT,
  507. timeout: httpx._types.TimeoutTypes | httpx._client.UseClientDefault = httpx._client.USE_CLIENT_DEFAULT,
  508. extensions: dict[str, Any] | None = None,
  509. ) -> httpx.Response:
  510. return super().put(
  511. url,
  512. content=content,
  513. data=data,
  514. files=files,
  515. json=json,
  516. params=params,
  517. headers=headers,
  518. cookies=cookies,
  519. auth=auth,
  520. follow_redirects=follow_redirects,
  521. timeout=timeout,
  522. extensions=extensions,
  523. )
  524. def patch( # type: ignore[override]
  525. self,
  526. url: httpx._types.URLTypes,
  527. *,
  528. content: httpx._types.RequestContent | None = None,
  529. data: _RequestData | None = None,
  530. files: httpx._types.RequestFiles | None = None,
  531. json: Any = None,
  532. params: httpx._types.QueryParamTypes | None = None,
  533. headers: httpx._types.HeaderTypes | None = None,
  534. cookies: httpx._types.CookieTypes | None = None,
  535. auth: httpx._types.AuthTypes | httpx._client.UseClientDefault = httpx._client.USE_CLIENT_DEFAULT,
  536. follow_redirects: bool | httpx._client.UseClientDefault = httpx._client.USE_CLIENT_DEFAULT,
  537. timeout: httpx._types.TimeoutTypes | httpx._client.UseClientDefault = httpx._client.USE_CLIENT_DEFAULT,
  538. extensions: dict[str, Any] | None = None,
  539. ) -> httpx.Response:
  540. return super().patch(
  541. url,
  542. content=content,
  543. data=data,
  544. files=files,
  545. json=json,
  546. params=params,
  547. headers=headers,
  548. cookies=cookies,
  549. auth=auth,
  550. follow_redirects=follow_redirects,
  551. timeout=timeout,
  552. extensions=extensions,
  553. )
  554. def delete( # type: ignore[override]
  555. self,
  556. url: httpx._types.URLTypes,
  557. *,
  558. params: httpx._types.QueryParamTypes | None = None,
  559. headers: httpx._types.HeaderTypes | None = None,
  560. cookies: httpx._types.CookieTypes | None = None,
  561. auth: httpx._types.AuthTypes | httpx._client.UseClientDefault = httpx._client.USE_CLIENT_DEFAULT,
  562. follow_redirects: bool | httpx._client.UseClientDefault = httpx._client.USE_CLIENT_DEFAULT,
  563. timeout: httpx._types.TimeoutTypes | httpx._client.UseClientDefault = httpx._client.USE_CLIENT_DEFAULT,
  564. extensions: dict[str, Any] | None = None,
  565. ) -> httpx.Response:
  566. return super().delete(
  567. url,
  568. params=params,
  569. headers=headers,
  570. cookies=cookies,
  571. auth=auth,
  572. follow_redirects=follow_redirects,
  573. timeout=timeout,
  574. extensions=extensions,
  575. )
  576. def websocket_connect(
  577. self,
  578. url: str,
  579. subprotocols: Sequence[str] | None = None,
  580. **kwargs: Any,
  581. ) -> WebSocketTestSession:
  582. url = urljoin("ws://testserver", url)
  583. headers = kwargs.get("headers", {})
  584. headers.setdefault("connection", "upgrade")
  585. headers.setdefault("sec-websocket-key", "testserver==")
  586. headers.setdefault("sec-websocket-version", "13")
  587. if subprotocols is not None:
  588. headers.setdefault("sec-websocket-protocol", ", ".join(subprotocols))
  589. kwargs["headers"] = headers
  590. try:
  591. super().request("GET", url, **kwargs)
  592. except _Upgrade as exc:
  593. session = exc.session
  594. else:
  595. raise RuntimeError("Expected WebSocket upgrade") # pragma: no cover
  596. return session
  597. def __enter__(self) -> Self:
  598. with contextlib.ExitStack() as stack:
  599. self.portal = portal = stack.enter_context(anyio.from_thread.start_blocking_portal(**self.async_backend))
  600. @stack.callback
  601. def reset_portal() -> None:
  602. self.portal = None
  603. send: anyio.create_memory_object_stream[MutableMapping[str, Any] | None] = (
  604. anyio.create_memory_object_stream(math.inf)
  605. )
  606. receive: anyio.create_memory_object_stream[MutableMapping[str, Any]] = anyio.create_memory_object_stream(
  607. math.inf
  608. )
  609. for channel in (*send, *receive):
  610. stack.callback(channel.close)
  611. self.stream_send = StapledObjectStream(*send)
  612. self.stream_receive = StapledObjectStream(*receive)
  613. self.task = portal.start_task_soon(self.lifespan)
  614. portal.call(self.wait_startup)
  615. @stack.callback
  616. def wait_shutdown() -> None:
  617. portal.call(self.wait_shutdown)
  618. self.exit_stack = stack.pop_all()
  619. return self
  620. def __exit__(self, *args: Any) -> None:
  621. self.exit_stack.close()
  622. async def lifespan(self) -> None:
  623. scope = {"type": "lifespan", "state": self.app_state}
  624. try:
  625. await self.app(scope, self.stream_receive.receive, self.stream_send.send)
  626. finally:
  627. await self.stream_send.send(None)
  628. async def wait_startup(self) -> None:
  629. await self.stream_receive.send({"type": "lifespan.startup"})
  630. async def receive() -> Any:
  631. message = await self.stream_send.receive()
  632. if message is None:
  633. self.task.result()
  634. return message
  635. message = await receive()
  636. assert message["type"] in (
  637. "lifespan.startup.complete",
  638. "lifespan.startup.failed",
  639. )
  640. if message["type"] == "lifespan.startup.failed":
  641. await receive()
  642. async def wait_shutdown(self) -> None:
  643. async def receive() -> Any:
  644. message = await self.stream_send.receive()
  645. if message is None:
  646. self.task.result()
  647. return message
  648. await self.stream_receive.send({"type": "lifespan.shutdown"})
  649. message = await receive()
  650. assert message["type"] in (
  651. "lifespan.shutdown.complete",
  652. "lifespan.shutdown.failed",
  653. )
  654. if message["type"] == "lifespan.shutdown.failed":
  655. await receive()