testclient.py 29 KB

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