testclient.py 29 KB

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