From 95a5e76230a166e05d9f03ce5304e8c32d7d1c96 Mon Sep 17 00:00:00 2001 From: Saurabh Date: Sun, 5 Jul 2026 00:46:25 +0530 Subject: [PATCH] =?UTF-8?q?=E2=9A=A1=EF=B8=8F=20Speed=20up=20stream=20canc?= =?UTF-8?q?ellation=20and=20security=20tutorial=20tests?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Cursor --- tests/test_stream_cancellation.py | 33 +++++++++++-------- .../test_security/test_tutorial004.py | 18 ++++++++++ 2 files changed, 37 insertions(+), 14 deletions(-) diff --git a/tests/test_stream_cancellation.py b/tests/test_stream_cancellation.py index 18e6d67d5..d15b1feb2 100644 --- a/tests/test_stream_cancellation.py +++ b/tests/test_stream_cancellation.py @@ -27,7 +27,7 @@ async def stream_raw() -> AsyncIterable[str]: i = 0 while True: yield f"item {i}\n" - i += 1 + i += 1 # pragma: no cover @app.get("/stream-jsonl") @@ -36,15 +36,20 @@ async def stream_jsonl() -> AsyncIterable[int]: i = 0 while True: yield i - i += 1 + i += 1 # pragma: no cover -async def _run_asgi_and_cancel(app: FastAPI, path: str, timeout: float) -> bool: - """Call the ASGI app for *path* and cancel after *timeout* seconds. +async def _run_asgi_and_cancel( + app: FastAPI, path: str, *, timeout: float = 10.0 +) -> bool: + """Call the ASGI app for *path* and cancel as soon as the first body chunk arrives. - Returns `True` if the cancellation was delivered (i.e. it did not hang). + Returns `True` if cancellation was delivered at a checkpoint (and the stream + actually produced data). A *timeout* safety net prevents the test from hanging + if cancellation never gets delivered. """ chunks: list[bytes] = [] + cancel_scope: anyio.CancelScope | None = None async def receive() -> Message: # Simulate a client that never disconnects, rely on cancellation @@ -52,8 +57,11 @@ async def _run_asgi_and_cancel(app: FastAPI, path: str, timeout: float) -> bool: return {"type": "http.disconnect"} # pragma: no cover async def send(message: Message) -> None: + nonlocal cancel_scope if message["type"] == "http.response.body": chunks.append(message.get("body", b"")) + if cancel_scope is not None and not cancel_scope.cancel_called: + cancel_scope.cancel() scope: Scope = { "type": "http", @@ -67,23 +75,20 @@ async def _run_asgi_and_cancel(app: FastAPI, path: str, timeout: float) -> bool: "server": ("test", 80), } - with anyio.move_on_after(timeout) as cancel_scope: - await app(scope, receive, send) + with anyio.move_on_after(timeout): + with anyio.CancelScope() as cancel_scope: + await app(scope, receive, send) - # If we got here within the timeout the generator was cancellable. - # cancel_scope.cancelled_caught is True when move_on_after fired. - return cancel_scope.cancelled_caught or len(chunks) > 0 + return cancel_scope.cancelled_caught and len(chunks) > 0 async def test_raw_stream_cancellation() -> None: """Raw streaming endpoint should be cancellable within a reasonable time.""" - cancelled = await _run_asgi_and_cancel(app, "/stream-raw", timeout=3.0) - # The key assertion: we reached this line at all (didn't hang). - # cancelled will be True because the infinite generator was interrupted. + cancelled = await _run_asgi_and_cancel(app, "/stream-raw") assert cancelled async def test_jsonl_stream_cancellation() -> None: """JSONL streaming endpoint should be cancellable within a reasonable time.""" - cancelled = await _run_asgi_and_cancel(app, "/stream-jsonl", timeout=3.0) + cancelled = await _run_asgi_and_cancel(app, "/stream-jsonl") assert cancelled diff --git a/tests/test_tutorial/test_security/test_tutorial004.py b/tests/test_tutorial/test_security/test_tutorial004.py index e52a029bd..b86ce38dd 100644 --- a/tests/test_tutorial/test_security/test_tutorial004.py +++ b/tests/test_tutorial/test_security/test_tutorial004.py @@ -1,5 +1,7 @@ import importlib +from functools import lru_cache from types import ModuleType +from typing import Any, cast from unittest.mock import patch import pytest @@ -15,6 +17,7 @@ from ...utils import needs_py310 pytest.param("tutorial004_py310", marks=needs_py310), pytest.param("tutorial004_an_py310", marks=needs_py310), ], + scope="module", ) def get_mod(request: pytest.FixtureRequest): mod = importlib.import_module(f"docs_src.security.{request.param}") @@ -22,6 +25,21 @@ def get_mod(request: pytest.FixtureRequest): return mod +@pytest.fixture(scope="module", autouse=True) +def cache_verify_password(mod: ModuleType): + assert hasattr(mod, "verify_password"), ( + f"Module {mod.__name__} does not have attribute 'verify_password'" + ) + + mod_any = cast(Any, mod) + original_func = mod_any.verify_password + cached_func = lru_cache()(original_func) + + mod_any.verify_password = cached_func + yield + mod_any.verify_password = original_func + + def get_access_token(*, username="johndoe", password="secret", client: TestClient): data = {"username": username, "password": password} response = client.post("/token", data=data)