Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/starlette/testclient.py: 21%

1from __future__ import annotations

3import contextlib

4import inspect

5import json

6import math

7import warnings

8from collections.abc import Awaitable, Callable, Generator, Iterable, Mapping, Sequence

9from concurrent.futures import Future, wait as wait_futures

10from contextlib import AbstractContextManager

11from types import GeneratorType

12from typing import TYPE_CHECKING, Any, Literal, Self, TypedDict, TypeGuard, cast

13from urllib.parse import unquote, urljoin

15import anyio

16import anyio.abc

17import anyio.from_thread

18from anyio.streams.stapled import StapledObjectStream

20from starlette._utils import is_async_callable

21from starlette.exceptions import StarletteDeprecationWarning

22from starlette.types import ASGIApp, Message, Receive, Scope, Send

23from starlette.websockets import WebSocketDisconnect

25if TYPE_CHECKING:

26 import httpx2 as httpx

27else:

28 try:

29 import httpx2 as httpx

30 except ModuleNotFoundError: # pragma: no cover

31 try:

32 import httpx

33 except ModuleNotFoundError:

34 raise RuntimeError(

35 "The starlette.testclient module requires the httpx2 package to be installed.\n"

36 "You can install this with:\n"

37 " $ pip install httpx2\n"

38 ) from None

39 else:

40 warnings.warn(

41 "Using `httpx` with `starlette.testclient` is deprecated; install `httpx2` instead.",

42 StarletteDeprecationWarning,

43 stacklevel=2,

44 )

46_PortalFactoryType = Callable[[], AbstractContextManager[anyio.from_thread.BlockingPortal]]

48ASGIInstance = Callable[[Receive, Send], Awaitable[None]]

49ASGI2App = Callable[[Scope], ASGIInstance]

50ASGI3App = Callable[[Scope, Receive, Send], Awaitable[None]]

53_RequestData = Mapping[str, str | Iterable[str] | bytes]

56def _is_asgi3(app: ASGI2App | ASGI3App) -> TypeGuard[ASGI3App]:

57 if inspect.isclass(app):

58 return hasattr(app, "__await__")

59 return is_async_callable(app)

62class _WrapASGI2:

63 """

64 Provide an ASGI3 interface onto an ASGI2 app.

65 """

67 def __init__(self, app: ASGI2App) -> None:

68 self.app = app

70 async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:

71 instance = self.app(scope)

72 await instance(receive, send)

75class _AsyncBackend(TypedDict):

76 backend: str

77 backend_options: dict[str, Any]

80class _Upgrade(Exception):

81 def __init__(self, session: WebSocketTestSession) -> None:

82 self.session = session

85class WebSocketDenialResponse( # type: ignore[misc]

86 httpx.Response,

87 WebSocketDisconnect,

88):

89 """

90 A special case of `WebSocketDisconnect`, raised in the `TestClient` if the

91 `WebSocket` is closed before being accepted with a `send_denial_response()`.

92 """

95class WebSocketTestSession:

96 def __init__(

97 self,

98 app: ASGI3App,

99 scope: Scope,

100 portal_factory: _PortalFactoryType,

101 ) -> None:

102 self.app = app

103 self.scope = scope

104 self.accepted_subprotocol = None

105 self.portal_factory = portal_factory

106 self.extra_headers = None

107

108 def __enter__(self) -> Self:

109 with contextlib.ExitStack() as stack:

110 self.portal = portal = stack.enter_context(self.portal_factory())

111 fut, cs = portal.start_task(self._run)

112 stack.callback(fut.result)

113 stack.callback(portal.call, cs.cancel)

114 self.send({"type": "websocket.connect"})

115 message = self.receive()

116 self._raise_on_close(message)

117 self.accepted_subprotocol = message.get("subprotocol", None)

118 self.extra_headers = message.get("headers", None)

119 stack.callback(self.close, 1000)

120 self.exit_stack = stack.pop_all()

121 return self

122

123 def __exit__(self, *args: Any) -> bool | None:

124 return self.exit_stack.__exit__(*args)

125

126 async def _run(self, *, task_status: anyio.abc.TaskStatus[anyio.CancelScope]) -> None:

127 """

128 The sub-thread in which the websocket session runs.

129 """

130 send_tx, send_rx = anyio.create_memory_object_stream[Message](math.inf)

131 receive_tx, receive_rx = anyio.create_memory_object_stream[Message](math.inf)

132 with send_tx, send_rx, receive_tx, receive_rx, anyio.CancelScope() as cs:

133 self._receive_tx = receive_tx

134 self._send_rx = send_rx

135 task_status.started(cs)

136 await self.app(self.scope, receive_rx.receive, send_tx.send)

137

138 # wait for cs.cancel to be called before closing streams

139 await anyio.sleep_forever()

140

141 def _raise_on_close(self, message: Message) -> None:

142 if message["type"] == "websocket.close":

143 raise WebSocketDisconnect(code=message.get("code", 1000), reason=message.get("reason", ""))

144 elif message["type"] == "websocket.http.response.start":

145 status_code: int = message["status"]

146 headers: list[tuple[bytes, bytes]] = message["headers"]

147 body: list[bytes] = []

148 while True:

149 message = self.receive()

150 assert message["type"] == "websocket.http.response.body"

151 body.append(message["body"])

152 if not message.get("more_body", False):

153 break

154 raise WebSocketDenialResponse(status_code=status_code, headers=headers, content=b"".join(body))

155

156 def send(self, message: Message) -> None:

157 self.portal.call(self._receive_tx.send, message)

158

159 def send_text(self, data: str) -> None:

160 self.send({"type": "websocket.receive", "text": data})

161

162 def send_bytes(self, data: bytes) -> None:

163 self.send({"type": "websocket.receive", "bytes": data})

164

165 def send_json(self, data: Any, mode: Literal["text", "binary"] = "text") -> None:

166 text = json.dumps(data, separators=(",", ":"), ensure_ascii=False)

167 if mode == "text":

168 self.send({"type": "websocket.receive", "text": text})

169 else:

170 self.send({"type": "websocket.receive", "bytes": text.encode("utf-8")})

171

172 def close(self, code: int = 1000, reason: str | None = None) -> None:

173 self.send({"type": "websocket.disconnect", "code": code, "reason": reason})

174

175 def receive(self) -> Message:

176 return self.portal.call(self._send_rx.receive)

177

178 def receive_text(self) -> str:

179 message = self.receive()

180 self._raise_on_close(message)

181 return cast(str, message["text"])

182

183 def receive_bytes(self) -> bytes:

184 message = self.receive()

185 self._raise_on_close(message)

186 return cast(bytes, message["bytes"])

187

188 def receive_json(self, mode: Literal["text", "binary"] = "text") -> Any:

189 message = self.receive()

190 self._raise_on_close(message)

191 if mode == "text":

192 text = message["text"]

193 else:

194 text = message["bytes"].decode("utf-8")

195 return json.loads(text)

196

197

198class _TestClientResponseStream(httpx.SyncByteStream):

199 def __init__(self, body: Generator[bytes, None, None], exit_stack: contextlib.ExitStack) -> None:

200 self._body = body

201 self._exit_stack = exit_stack

202

203 def __iter__(self) -> Generator[bytes, None, None]:

204 return self._body

205

206 def close(self) -> None:

207 self._exit_stack.close()

208

209

210class _TestClientTransport(httpx.BaseTransport):

211 def __init__(

212 self,

213 app: ASGI3App,

214 portal_factory: _PortalFactoryType,

215 raise_server_exceptions: bool = True,

216 root_path: str = "",

217 *,

218 client: tuple[str, int],

219 app_state: dict[str, Any],

220 ) -> None:

221 self.app = app

222 self.raise_server_exceptions = raise_server_exceptions

223 self.root_path = root_path

224 self.portal_factory = portal_factory

225 self.app_state = app_state

226 self.client = client

227 self._open_requests: set[Callable[[], None]] = set()

228

229 def close(self) -> None:

230 with contextlib.ExitStack() as stack:

231 while self._open_requests:

232 stack.callback(self._open_requests.pop())

233

234 def handle_request(self, request: httpx.Request) -> httpx.Response:

235 scheme = request.url.scheme

236 host = request.url.raw_host.decode(encoding="ascii")

237 path = request.url.path

238 raw_path = request.url.raw_path

239 query = request.url.query.decode(encoding="ascii")

240

241 default_port = {"http": 80, "ws": 80, "https": 443, "wss": 443}[scheme]

242 port = request.url.port

243 if port is None:

244 port = default_port

245

246 # Include the 'host' header.

247 if "host" in request.headers:

248 headers: list[tuple[bytes, bytes]] = []

249 else: # pragma: no cover

250 headers = [(b"host", request.url.netloc)]

251

252 # Include other request headers.

253 headers += [(key.lower().encode(), value.encode()) for key, value in request.headers.multi_items()]

254

255 scope: dict[str, Any]

256

257 if scheme in {"ws", "wss"}:

258 subprotocol = request.headers.get("sec-websocket-protocol", None)

259 if subprotocol is None:

260 subprotocols: Sequence[str] = []

261 else:

262 subprotocols = [value.strip() for value in subprotocol.split(",")]

263 scope = {

264 "type": "websocket",

265 "path": unquote(path),

266 "raw_path": raw_path.split(b"?", 1)[0],

267 "root_path": self.root_path,

268 "scheme": scheme,

269 "query_string": query.encode(),

270 "headers": headers,

271 "client": self.client,

272 "server": [host, port],

273 "subprotocols": subprotocols,

274 "state": self.app_state.copy(),

275 "extensions": {"websocket.http.response": {}},

276 }

277 session = WebSocketTestSession(self.app, scope, self.portal_factory)

278 raise _Upgrade(session)

279

280 scope = {

281 "type": "http",

282 "http_version": "1.1",

283 "method": request.method,

284 "path": unquote(path),

285 "raw_path": raw_path.split(b"?", 1)[0],

286 "root_path": self.root_path,

287 "scheme": scheme,

288 "query_string": query.encode(),

289 "headers": headers,

290 "client": self.client,

291 "server": [host, port],

292 "extensions": {"http.response.debug": {}, "http.response.trailers": {}},

293 "state": self.app_state.copy(),

294 }

295

296 request_complete = False

297 response_started = False

298 body_complete = False

299 trailers_expected = False

300 trailers: list[tuple[bytes, bytes]] = []

301 response_complete: anyio.Event

302 raw_kwargs: dict[str, Any] = {}

303 debug_info: dict[str, Any] | None = None

304

305 async def receive() -> Message:

306 nonlocal request_complete

307

308 if request_complete:

309 if not response_complete.is_set():

310 await response_complete.wait()

311 return {"type": "http.disconnect"}

312

313 body = request.read()

314 if isinstance(body, str):

315 body_bytes: bytes = body.encode("utf-8") # pragma: no cover

316 elif body is None:

317 body_bytes = b"" # pragma: no cover

318 elif isinstance(body, GeneratorType):

319 try: # pragma: no cover

320 chunk = body.send(None)

321 if isinstance(chunk, str):

322 chunk = chunk.encode("utf-8")

323 return {"type": "http.request", "body": chunk, "more_body": True}

324 except StopIteration: # pragma: no cover

325 request_complete = True

326 return {"type": "http.request", "body": b""}

327 else:

328 body_bytes = body

329

330 request_complete = True

331 return {"type": "http.request", "body": body_bytes}

332

333 async def send(message: Message) -> None:

334 nonlocal raw_kwargs, response_started, debug_info, body_complete, trailers_expected

335

336 if message["type"] == "http.response.start":

337 assert not response_started, 'Received multiple "http.response.start" messages.'

338 raw_kwargs["status_code"] = message["status"]

339 raw_kwargs["headers"] = [(key.decode(), value.decode()) for key, value in message.get("headers", [])]

340 response_started = True

341 trailers_expected = message.get("trailers", False)

342 response_available.set()

343 elif message["type"] == "http.response.body":

344 assert response_started, 'Received "http.response.body" without "http.response.start".'

345 assert not response_complete.is_set(), 'Received "http.response.body" after response completed.'

346 assert not body_complete, 'Received "http.response.body" after body completed.'

347 body = message.get("body", b"")

348 more_body = message.get("more_body", False)

349 await body_tx.send((bytes(body), not more_body and not trailers_expected))

350 if not more_body:

351 body_complete = True

352 elif message["type"] == "http.response.trailers":

353 assert trailers_expected, 'Received "http.response.trailers" without declaring trailers.'

354 assert body_complete, 'Received "http.response.trailers" before body completed.'

355 assert not response_complete.is_set(), 'Received "http.response.trailers" after response completed.'

356 trailers.extend(message.get("headers", []))

357 if not message.get("more_trailers", False):

358 response_complete.set()

359 elif message["type"] == "http.response.debug":

360 debug_info = message["info"]

361

362 async def run_app() -> None:

363 try:

364 with cancel_scope:

365 async with body_tx:

366 await self.app(scope, receive, send)

367 finally:

368 response_available.set()

369

370 async def receive_body() -> bytes:

371 body, complete = await body_rx.receive()

372 if complete:

373 response_complete.set()

374 return body

375

376 def iter_body() -> Generator[bytes, None, None]:

377 while True:

378 try:

379 body = portal.call(receive_body)

380 except anyio.EndOfStream:

381 break

382 if request.method != "HEAD" and body:

383 yield body

384 exit_stack.close()

385

386 def cancel_if_incomplete() -> None:

387 if not response_complete.is_set():

388 cancel_scope.cancel()

389

390 def close_app(exc_type: type[BaseException] | None, *_: Any) -> None:

391 if app_task is None: # pragma: no cover - only if a signal interrupts start_task_soon()

392 portal.call(cancel_scope.cancel)

393 return

394 if not app_task.done():

395 portal.call(cancel_if_incomplete)

396 wait_futures([app_task])

397 self._open_requests.discard(abort)

398 if exc_type is None and self.raise_server_exceptions and (exception := app_task.exception()):

399 raise exception

400

401 def abort() -> None:

402 portal.call(cancel_scope.cancel)

403 exit_stack.close()

404

405 with contextlib.ExitStack() as stack:

406 portal = stack.enter_context(self.portal_factory())

407 response_available = portal.call(anyio.Event)

408 response_complete = portal.call(anyio.Event)

409 body_tx, body_rx = portal.call(anyio.create_memory_object_stream[tuple[bytes, bool]], 0)

410 stack.callback(body_tx.close)

411 stack.callback(body_rx.close)

412 cancel_scope = portal.call(anyio.CancelScope)

413 app_task: Future[None] | None = None

414 stack.push(close_app)

415 app_task = portal.start_task_soon(run_app)

416 portal.call(response_available.wait)

417

418 if not response_started:

419 stack.close()

420 if self.raise_server_exceptions:

421 raise AssertionError("TestClient did not receive any response.")

422 raw_kwargs = {"status_code": 500, "headers": [], "stream": httpx.ByteStream(b"")}

423 else:

424 exit_stack = stack.pop_all()

425 self._open_requests.add(abort)

426 raw_kwargs["stream"] = _TestClientResponseStream(iter_body(), exit_stack)

427

428 response = httpx.Response(**raw_kwargs, request=request)

429 if trailers_expected:

430 response.extensions["http.response.trailers"] = trailers

431 if debug_info is not None:

432 response.extensions["http.response.debug"] = debug_info

433 if "template" in debug_info:

434 response.template = debug_info["template"] # type: ignore[attr-defined]

435 if "context" in debug_info:

436 response.context = debug_info["context"] # type: ignore[attr-defined]

437 return response

438

439

440class TestClient(httpx.Client):

441 __test__ = False

442 task: Future[None]

443 portal: anyio.from_thread.BlockingPortal | None = None

444

445 def __init__(

446 self,

447 app: ASGIApp,

448 base_url: str = "http://testserver",

449 raise_server_exceptions: bool = True,

450 root_path: str = "",

451 backend: Literal["asyncio", "trio"] = "asyncio",

452 backend_options: dict[str, Any] | None = None,

453 cookies: httpx._types.CookieTypes | None = None,

454 headers: dict[str, str] | None = None,

455 follow_redirects: bool = True,

456 client: tuple[str, int] = ("testclient", 50000),

457 ) -> None:

458 self.async_backend = _AsyncBackend(backend=backend, backend_options=backend_options or {})

459 if _is_asgi3(app):

460 asgi_app = app

461 else:

462 app = cast(ASGI2App, app) # type: ignore[assignment]

463 asgi_app = _WrapASGI2(app) # type: ignore[arg-type]

464 self.app = asgi_app

465 self.app_state: dict[str, Any] = {}

466 transport = _TestClientTransport(

467 self.app,

468 portal_factory=self._portal_factory,

469 raise_server_exceptions=raise_server_exceptions,

470 root_path=root_path,

471 app_state=self.app_state,

472 client=client,

473 )

474 if headers is None:

475 headers = {}

476 headers.setdefault("user-agent", "testclient")

477 super().__init__(

478 base_url=base_url,

479 headers=headers,

480 transport=transport,

481 follow_redirects=follow_redirects,

482 cookies=cookies,

483 )

484

485 @contextlib.contextmanager

486 def _portal_factory(self) -> Generator[anyio.from_thread.BlockingPortal, None, None]:

487 if self.portal is not None:

488 yield self.portal

489 else:

490 with anyio.from_thread.start_blocking_portal(**self.async_backend) as portal:

491 yield portal

492

493 def request( # type: ignore[override]

494 self,

495 method: str,

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 if timeout is not httpx.USE_CLIENT_DEFAULT:

511 warnings.warn(

512 "You should not use the 'timeout' argument with the TestClient. "

513 "See https://github.com/Kludex/starlette/issues/1108 for more information.",

514 StarletteDeprecationWarning,

515 stacklevel=2,

516 )

517 url = self._merge_url(url)

518 return super().request(

519 method,

520 url,

521 content=content,

522 data=data,

523 files=files,

524 json=json,

525 params=params,

526 headers=headers,

527 cookies=cookies,

528 auth=auth,

529 follow_redirects=follow_redirects,

530 timeout=timeout,

531 extensions=extensions,

532 )

533

534 def websocket_connect(

535 self,

536 url: str,

537 subprotocols: Sequence[str] | None = None,

538 **kwargs: Any,

539 ) -> WebSocketTestSession:

540 url = urljoin("ws://testserver", url)

541 headers = kwargs.get("headers", {})

542 headers.setdefault("connection", "upgrade")

543 headers.setdefault("sec-websocket-key", "testserver==")

544 headers.setdefault("sec-websocket-version", "13")

545 if subprotocols is not None:

546 headers.setdefault("sec-websocket-protocol", ", ".join(subprotocols))

547 kwargs["headers"] = headers

548 try:

549 super().request("GET", url, **kwargs)

550 except _Upgrade as exc:

551 session = exc.session

552 else:

553 raise RuntimeError("Expected WebSocket upgrade") # pragma: no cover

554

555 return session

556

557 def __enter__(self) -> Self:

558 with contextlib.ExitStack() as stack:

559 self.portal = portal = stack.enter_context(anyio.from_thread.start_blocking_portal(**self.async_backend))

560

561 @stack.callback

562 def reset_portal() -> None:

563 self.portal = None

564

565 send = anyio.create_memory_object_stream[Message | None](math.inf)

566 receive = anyio.create_memory_object_stream[Message](math.inf)

567 for channel in (*send, *receive):

568 stack.callback(channel.close)

569 self.stream_send = StapledObjectStream(*send)

570 self.stream_receive = StapledObjectStream(*receive)

571 self.task = portal.start_task_soon(self.lifespan)

572 portal.call(self.wait_startup)

573

574 @stack.callback

575 def wait_shutdown() -> None:

576 portal.call(self.wait_shutdown)

577

578 stack.callback(self._transport.close)

579 self.exit_stack = stack.pop_all()

580

581 return self

582

583 def __exit__(self, *args: Any) -> None:

584 self.exit_stack.close()

585

586 async def lifespan(self) -> None:

587 scope = {"type": "lifespan", "state": self.app_state}

588 try:

589 await self.app(scope, self.stream_receive.receive, self.stream_send.send)

590 finally:

591 await self.stream_send.send(None)

592

593 async def wait_startup(self) -> None:

594 await self.stream_receive.send({"type": "lifespan.startup"})

595

596 async def receive() -> Any:

597 message = await self.stream_send.receive()

598 if message is None:

599 self.task.result()

600 return message

601

602 message = await receive()

603 assert message["type"] in (

604 "lifespan.startup.complete",

605 "lifespan.startup.failed",

606 )

607 if message["type"] == "lifespan.startup.failed":

608 await receive()

609

610 async def wait_shutdown(self) -> None:

611 async def receive() -> Any:

612 message = await self.stream_send.receive()

613 if message is None:

614 self.task.result()

615 return message

616

617 await self.stream_receive.send({"type": "lifespan.shutdown"})

618 message = await receive()

619 assert message["type"] in (

620 "lifespan.shutdown.complete",

621 "lifespan.shutdown.failed",

622 )

623 if message["type"] == "lifespan.shutdown.failed":

624 await receive()