Coverage for tests/test_dependency_after_yield_websockets.py: 100%

1from contextlib import contextmanager 7 ctx1abcdefg

2from typing import Any, Generator 7 ctx1abcdefg

3 

4import pytest 7 ctx1abcdefg

5from fastapi import Depends, FastAPI, WebSocket 7 ctx1abcdefg

6from fastapi.testclient import TestClient 7 ctx1abcdefg

7from typing_extensions import Annotated 7 ctx1abcdefg

8 

9 

10class Session: 7 ctx1abcdefg

11 def __init__(self) -> None: 7 ctx1abcdefg

12 self.data = ["foo", "bar", "baz"] 14 ctx1hoipjqkrlsmtnu

13 self.open = True 14 ctx1hoipjqkrlsmtnu

14 

15 def __iter__(self) -> Generator[str, None, None]: 7 ctx1abcdefg

16 for item in self.data: 14 ctx1hoipjqkrlsmtnu

17 if self.open: 14 ctx1hoipjqkrlsmtnu

18 yield item 7 ctx1hijklmn

19 else: 

20 raise ValueError("Session closed") 7 ctx1opqrstu

21 

22 

23@contextmanager 7 ctx1abcdefg

24def acquire_session() -> Generator[Session, None, None]: 7 ctx1abcdefg

25 session = Session() 14 ctx1hoipjqkrlsmtnu

26 try: 14 ctx1hoipjqkrlsmtnu

27 yield session 14 ctx1hoipjqkrlsmtnu

28 finally: 

29 session.open = False 14 ctx1hoipjqkrlsmtnu

30 

31 

32def dep_session() -> Any: 7 ctx1abcdefg

33 with acquire_session() as s: 7 ctx1hijklmn

34 yield s 7 ctx1hijklmn

35 

36 

37def broken_dep_session() -> Any: 7 ctx1abcdefg

38 with acquire_session() as s: 7 ctx1opqrstu

39 s.open = False 7 ctx1opqrstu

40 yield s 7 ctx1opqrstu

41 

42 

43SessionDep = Annotated[Session, Depends(dep_session)] 7 ctx1abcdefg

44BrokenSessionDep = Annotated[Session, Depends(broken_dep_session)] 7 ctx1abcdefg

45 

46app = FastAPI() 7 ctx1abcdefg

47 

48 

49@app.websocket("/ws") 7 ctx1abcdefg

50async def websocket_endpoint(websocket: WebSocket, session: SessionDep): 7 ctx1abcdefg

51 await websocket.accept() 7 ctx1hijklmn

52 for item in session: 7 ctx1hijklmn

53 await websocket.send_text(f"{item}") 7 ctx1hijklmn

54 

55 

56@app.websocket("/ws-broken") 7 ctx1abcdefg

57async def websocket_endpoint_broken(websocket: WebSocket, session: BrokenSessionDep): 7 ctx1abcdefg

58 await websocket.accept() 7 ctx1opqrstu

59 for item in session: 7 ctx1opqrstu

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

61 

62 

63client = TestClient(app) 7 ctx1abcdefg

64 

65 

66def test_websocket_dependency_after_yield(): 7 ctx1abcdefg

67 with client.websocket_connect("/ws") as websocket: 7 ctx1hijklmn

68 data = websocket.receive_text() 7 ctx1hijklmn

69 assert data == "foo" 7 ctx1hijklmn

70 data = websocket.receive_text() 7 ctx1hijklmn

71 assert data == "bar" 7 ctx1hijklmn

72 data = websocket.receive_text() 7 ctx1hijklmn

73 assert data == "baz" 7 ctx1hijklmn

74 

75 

76def test_websocket_dependency_after_yield_broken(): 7 ctx1abcdefg

77 with pytest.raises(ValueError, match="Session closed"): 7 ctx1opqrstu

78 with client.websocket_connect("/ws-broken"): 7 ctx1opqrstu

79 pass # pragma no cover 7 ctx1opqrstu