Coverage for tests/test_dependency_contextvars.py: 100%

1from contextvars import ContextVar 6 ctx1abcdef

2from typing import Any, Awaitable, Callable, Dict, Optional 6 ctx1abcdef

3 

4from fastapi import Depends, FastAPI, Request, Response 6 ctx1abcdef

5from fastapi.testclient import TestClient 6 ctx1abcdef

6 

7legacy_request_state_context_var: ContextVar[Optional[Dict[str, Any]]] = ContextVar( 6 ctx1abcdef

8 "legacy_request_state_context_var", default=None 

9) 

10 

11app = FastAPI() 6 ctx1abcdef

12 

13 

14async def set_up_request_state_dependency(): 6 ctx1abcdef

15 request_state = {"user": "deadpond"} 6 ctx1ghijkl

16 contextvar_token = legacy_request_state_context_var.set(request_state) 6 ctx1ghijkl

17 yield request_state 6 ctx1ghijkl

18 legacy_request_state_context_var.reset(contextvar_token) 6 ctx1ghijkl

19 

20 

21@app.middleware("http") 6 ctx1abcdef

22async def custom_middleware( 6 ctx1abcdef

23 request: Request, call_next: Callable[[Request], Awaitable[Response]] 

24): 

25 response = await call_next(request) 6 ctx1ghijkl

26 response.headers["custom"] = "foo" 6 ctx1ghijkl

27 return response 6 ctx1ghijkl

28 

29 

30@app.get("/user", dependencies=[Depends(set_up_request_state_dependency)]) 6 ctx1abcdef

31def get_user(): 6 ctx1abcdef

32 request_state = legacy_request_state_context_var.get() 6 ctx1ghijkl

33 assert request_state 6 ctx1ghijkl

34 return request_state["user"] 6 ctx1ghijkl

35 

36 

37client = TestClient(app) 6 ctx1abcdef

38 

39 

40def test_dependency_contextvars(): 6 ctx1abcdef

41 """ 

42 Check that custom middlewares don't affect the contextvar context for dependencies. 

43 

44 The code before yield and the code after yield should be run in the same contextvar 

45 context, so that request_state_context_var.reset(contextvar_token). 

46 

47 If they are run in a different context, that raises an error. 

48 """ 

49 response = client.get("/user") 6 ctx1ghijkl

50 assert response.json() == "deadpond" 6 ctx1ghijkl

51 assert response.headers["custom"] == "foo" 6 ctx1ghijkl