diff --git a/tests/test_depends_limiter.py b/tests/test_depends_limiter.py index f7e015b7e..9158a3bb1 100644 --- a/tests/test_depends_limiter.py +++ b/tests/test_depends_limiter.py @@ -1,18 +1,15 @@ from contextlib import asynccontextmanager -import threading +import anyio +import sniffio from fastapi import Depends, FastAPI from fastapi.testclient import TestClient from typing_extensions import Annotated -import anyio -import sniffio def get_borrowed_tokens(): - return { - lname: l.borrowed_tokens - for lname, l in sorted(limiters.items()) - } + return {lname: l.borrowed_tokens for lname, l in sorted(limiters.items())} + def limited_dep(): # run this in the event loop thread: @@ -22,11 +19,12 @@ def limited_dep(): def init_limiters(): return { - 'default': None, # should be set in Lifespan handler - 'a': anyio.CapacityLimiter(5), - 'b': anyio.CapacityLimiter(3), + "default": None, # should be set in Lifespan handler + "a": anyio.CapacityLimiter(5), + "b": anyio.CapacityLimiter(3), } + # Note: # initializing CapacityLimiters at module level before the event loop has started # needs anyio >= 4.2.0; starlette currently requires anyio >= 3.4.0 @@ -39,16 +37,18 @@ except sniffio.AsyncLibraryNotFoundError: # not in an async context yet async def _init_limiters(): return init_limiters() + limiters = anyio.run(_init_limiters) @asynccontextmanager async def lifespan(app: FastAPI): - limiters['default'] = anyio.to_thread.current_default_thread_limiter() + limiters["default"] = anyio.to_thread.current_default_thread_limiter() yield { - 'limiters': limiters, + "limiters": limiters, } + app = FastAPI(lifespan=lifespan) @@ -56,70 +56,75 @@ app = FastAPI(lifespan=lifespan) async def root(borrowed_tokens: Annotated[dict, Depends(limited_dep)]): return borrowed_tokens + @app.get("/a") -async def a(borrowed_tokens: Annotated[dict, Depends(limited_dep, limiter=limiters['a'])]): +async def a( + borrowed_tokens: Annotated[dict, Depends(limited_dep, limiter=limiters["a"])], +): return borrowed_tokens + @app.get("/b") -async def b(borrowed_tokens: Annotated[dict, Depends(limited_dep, limiter=limiters['b'])]): +async def b( + borrowed_tokens: Annotated[dict, Depends(limited_dep, limiter=limiters["b"])], +): return borrowed_tokens - def test_depends_limiter(): with TestClient(app) as client: response = client.get("/") assert response.status_code == 200, response.text - assert response.json() == {"a":0,"b":0,"default":1} + assert response.json() == {"a": 0, "b": 0, "default": 1} response = client.get("/a") assert response.status_code == 200, response.text - assert response.json() == {"a":1,"b":0,"default":0} + assert response.json() == {"a": 1, "b": 0, "default": 0} response = client.get("/b") assert response.status_code == 200, response.text - assert response.json() == {"a":0,"b":1,"default":0} + assert response.json() == {"a": 0, "b": 1, "default": 0} def test_openapi_schema(): with TestClient(app) as client: response = client.get("/openapi.json") assert response.status_code == 200, response.text - assert response.json()['paths'] == { - '/': { - 'get': { - 'summary': 'Root', - 'operationId': 'root__get', - 'responses': { - '200': { - 'description': 'Successful Response', - 'content': {'application/json': {'schema': {}}} + assert response.json()["paths"] == { + "/": { + "get": { + "summary": "Root", + "operationId": "root__get", + "responses": { + "200": { + "description": "Successful Response", + "content": {"application/json": {"schema": {}}}, } - } + }, } }, - '/a': { - 'get': { - 'summary': 'A', - 'operationId': 'a_a_get', - 'responses': { - '200': { - 'description': 'Successful Response', - 'content': {'application/json': {'schema': {}}} + "/a": { + "get": { + "summary": "A", + "operationId": "a_a_get", + "responses": { + "200": { + "description": "Successful Response", + "content": {"application/json": {"schema": {}}}, } - } + }, } }, - '/b': { - 'get': { - 'summary': 'B', - 'operationId': 'b_b_get', - 'responses': { - '200': { - 'description': 'Successful Response', - 'content': {'application/json': {'schema': {}}} + "/b": { + "get": { + "summary": "B", + "operationId": "b_b_get", + "responses": { + "200": { + "description": "Successful Response", + "content": {"application/json": {"schema": {}}}, } - } + }, } - } + }, }