Coverage for tests/test_dependency_after_yield_websockets.py: 100%

1from contextlib import contextmanager 6 ctx1abcdef

2from typing import Any, Generator 6 ctx1abcdef

3 

4import pytest 6 ctx1abcdef

5from fastapi import Depends, FastAPI, WebSocket 6 ctx1abcdef

6from fastapi.testclient import TestClient 6 ctx1abcdef

7from typing_extensions import Annotated 6 ctx1abcdef

8 

9 

10class Session: 6 ctx1abcdef

11 def __init__(self) -> None: 6 ctx1abcdef

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

13 self.open = True 12 ctx1gmhniojpkqlr

14 

15 def __iter__(self) -> Generator[str, None, None]: 6 ctx1abcdef

16 for item in self.data: 12 ctx1gmhniojpkqlr

17 if self.open: 12 ctx1gmhniojpkqlr

18 yield item 6 ctx1ghijkl

19 else: 

20 raise ValueError("Session closed") 6 ctx1mnopqr

21 

22 

23@contextmanager 6 ctx1abcdef

24def acquire_session() -> Generator[Session, None, None]: 6 ctx1abcdef

25 session = Session() 12 ctx1gmhniojpkqlr

26 try: 12 ctx1gmhniojpkqlr

27 yield session 12 ctx1gmhniojpkqlr

28 finally: 

29 session.open = False 12 ctx1gmhniojpkqlr

30 

31 

32def dep_session() -> Any: 6 ctx1abcdef

33 with acquire_session() as s: 6 ctx1ghijkl

34 yield s 6 ctx1ghijkl

35 

36 

37def broken_dep_session() -> Any: 6 ctx1abcdef

38 with acquire_session() as s: 6 ctx1mnopqr

39 s.open = False 6 ctx1mnopqr

40 yield s 6 ctx1mnopqr

41 

42 

43SessionDep = Annotated[Session, Depends(dep_session)] 6 ctx1abcdef

44BrokenSessionDep = Annotated[Session, Depends(broken_dep_session)] 6 ctx1abcdef

45 

46app = FastAPI() 6 ctx1abcdef

47 

48 

49@app.websocket("/ws") 6 ctx1abcdef

50async def websocket_endpoint(websocket: WebSocket, session: SessionDep): 6 ctx1abcdef

51 await websocket.accept() 6 ctx1ghijkl

52 for item in session: 6 ctx1ghijkl

53 await websocket.send_text(f"{item}") 6 ctx1ghijkl

54 

55 

56@app.websocket("/ws-broken") 6 ctx1abcdef

57async def websocket_endpoint_broken(websocket: WebSocket, session: BrokenSessionDep): 6 ctx1abcdef

58 await websocket.accept() 6 ctx1mnopqr

59 for item in session: 6 ctx1mnopqr

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

61 

62 

63client = TestClient(app) 6 ctx1abcdef

64 

65 

66def test_websocket_dependency_after_yield(): 6 ctx1abcdef

67 with client.websocket_connect("/ws") as websocket: 6 ctx1ghijkl

68 data = websocket.receive_text() 6 ctx1ghijkl

69 assert data == "foo" 6 ctx1ghijkl

70 data = websocket.receive_text() 6 ctx1ghijkl

71 assert data == "bar" 6 ctx1ghijkl

72 data = websocket.receive_text() 6 ctx1ghijkl

73 assert data == "baz" 6 ctx1ghijkl

74 

75 

76def test_websocket_dependency_after_yield_broken(): 6 ctx1abcdef

77 with pytest.raises(ValueError, match="Session closed"): 6 ctx1mnopqr

78 with client.websocket_connect("/ws-broken"): 6 ctx1mnopqr

79 pass # pragma no cover 6 ctx1mnopqr