Coverage for tests / test_dependency_after_yield_websockets.py: 100%

1from collections.abc import Generator 4 ctx1abcd

2from contextlib import contextmanager 4 ctx1abcd

3from typing import Annotated, Any 4 ctx1abcd

4 

5import pytest 4 ctx1abcd

6from fastapi import Depends, FastAPI, WebSocket 4 ctx1abcd

7from fastapi.testclient import TestClient 4 ctx1abcd

8 

9 

10class Session: 4 ctx1abcd

11 def __init__(self) -> None: 4 ctx1abcd

12 self.data = ["foo", "bar", "baz"] 6 ctx1ejfhgi

13 self.open = True 6 ctx1ejfhgi

14 

15 def __iter__(self) -> Generator[str, None, None]: 4 ctx1abcd

16 for item in self.data: 6 ctx1ejfhgi

17 if self.open: 6 ctx1ejfhgi

18 yield item 3 ctx1efg

19 else: 

20 raise ValueError("Session closed") 3 ctx1jhi

21 

22 

23@contextmanager 4 ctx1abcd

24def acquire_session() -> Generator[Session, None, None]: 4 ctx1abcd

25 session = Session() 6 ctx1ejfhgi

26 try: 6 ctx1ejfhgi

27 yield session 6 ctx1ejfhgi

28 finally: 

29 session.open = False 6 ctx1ejfhgi

30 

31 

32def dep_session() -> Any: 4 ctx1abcd

33 with acquire_session() as s: 3 ctx1efg

34 yield s 3 ctx1efg

35 

36 

37def broken_dep_session() -> Any: 4 ctx1abcd

38 with acquire_session() as s: 3 ctx1jhi

39 s.open = False 3 ctx1jhi

40 yield s 3 ctx1jhi

41 

42 

43SessionDep = Annotated[Session, Depends(dep_session)] 4 ctx1abcd

44BrokenSessionDep = Annotated[Session, Depends(broken_dep_session)] 4 ctx1abcd

45 

46app = FastAPI() 4 ctx1abcd

47 

48 

49@app.websocket("/ws") 4 ctx1abcd

50async def websocket_endpoint(websocket: WebSocket, session: SessionDep): 4 ctx1abcd

51 await websocket.accept() 3 ctx1efg

52 for item in session: 3 ctx1efg

53 await websocket.send_text(f"{item}") 3 ctx1efg

54 

55 

56@app.websocket("/ws-broken") 4 ctx1abcd

57async def websocket_endpoint_broken(websocket: WebSocket, session: BrokenSessionDep): 4 ctx1abcd

58 await websocket.accept() 3 ctx1jhi

59 for item in session: 3 ctx1jhi

60 await websocket.send_text(f"{item}") # pragma no cover 

61 

62 

63client = TestClient(app) 4 ctx1abcd

64 

65 

66def test_websocket_dependency_after_yield(): 4 ctx1abcd

67 with client.websocket_connect("/ws") as websocket: 3 ctx1efg

68 data = websocket.receive_text() 3 ctx1efg

69 assert data == "foo" 3 ctx1efg

70 data = websocket.receive_text() 3 ctx1efg

71 assert data == "bar" 3 ctx1efg

72 data = websocket.receive_text() 3 ctx1efg

73 assert data == "baz" 3 ctx1efg

74 

75 

76def test_websocket_dependency_after_yield_broken(): 4 ctx1abcd

77 with pytest.raises(ValueError, match="Session closed"): 3 ctx1jhi

78 with client.websocket_connect("/ws-broken"): 3 ctx1jhi

79 pass # pragma no cover 2 ctx1hi