Coverage for tests/test_dependency_contextvars.py: 100%

1from contextvars import ContextVar 5 ctx1abcde

2from typing import Any, Awaitable, Callable, Dict, Optional 5 ctx1abcde

3 

4from fastapi import Depends, FastAPI, Request, Response 5 ctx1abcde

5from fastapi.testclient import TestClient 5 ctx1abcde

6 

7legacy_request_state_context_var: ContextVar[Optional[Dict[str, Any]]] = ContextVar( 5 ctx1abcde

8 "legacy_request_state_context_var", default=None 

9) 

10 

11app = FastAPI() 5 ctx1abcde

12 

13 

14async def set_up_request_state_dependency(): 5 ctx1abcde

15 request_state = {"user": "deadpond"} 5 ctx1fghij

16 contextvar_token = legacy_request_state_context_var.set(request_state) 5 ctx1fghij

17 yield request_state 5 ctx1fghij

18 legacy_request_state_context_var.reset(contextvar_token) 5 ctx1fghij

19 

20 

21@app.middleware("http") 5 ctx1abcde

22async def custom_middleware( 5 ctx1abcde

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

24): 

25 response = await call_next(request) 5 ctx1fghij

26 response.headers["custom"] = "foo" 5 ctx1fghij

27 return response 5 ctx1fghij

28 

29 

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

31def get_user(): 5 ctx1abcde

32 request_state = legacy_request_state_context_var.get() 5 ctx1fghij

33 assert request_state 5 ctx1fghij

34 return request_state["user"] 5 ctx1fghij

35 

36 

37client = TestClient(app) 5 ctx1abcde

38 

39 

40def test_dependency_contextvars(): 5 ctx1abcde

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") 5 ctx1fghij

50 assert response.json() == "deadpond" 5 ctx1fghij

51 assert response.headers["custom"] == "foo" 5 ctx1fghij