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 ctx1abcde

16 contextvar_token = legacy_request_state_context_var.set(request_state) 5 ctx1abcde

17 yield request_state 5 ctx1abcde

18 legacy_request_state_context_var.reset(contextvar_token) 5 ctx1abcde

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 ctx1abcde

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

27 return response 5 ctx1abcde

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 ctx1abcde

33 assert request_state 5 ctx1abcde

34 return request_state["user"] 5 ctx1abcde

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 ctx1abcde

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

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