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"] 8 ctx1ielfjgkh

13 self.open = True 8 ctx1ielfjgkh

14 

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

16 for item in self.data: 8 ctx1ielfjgkh

17 if self.open: 8 ctx1ielfjgkh

18 yield item 4 ctx1efgh

19 else: 

20 raise ValueError("Session closed") 4 ctx1iljk

21 

22 

23@contextmanager 4 ctx1abcd

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

25 session = Session() 8 ctx1ielfjgkh

26 try: 8 ctx1ielfjgkh

27 yield session 8 ctx1ielfjgkh

28 finally: 

29 session.open = False 8 ctx1ielfjgkh

30 

31 

32def dep_session() -> Any: 4 ctx1abcd

33 with acquire_session() as s: 4 ctx1efgh

34 yield s 4 ctx1efgh

35 

36 

37def broken_dep_session() -> Any: 4 ctx1abcd

38 with acquire_session() as s: 4 ctx1iljk

39 s.open = False 4 ctx1iljk

40 yield s 4 ctx1iljk

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() 4 ctx1efgh

52 for item in session: 4 ctx1efgh

53 await websocket.send_text(f"{item}") 4 ctx1efgh

54 

55 

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

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

58 await websocket.accept() 4 ctx1iljk

59 for item in session: 4 ctx1iljk

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: 4 ctx1efgh

68 data = websocket.receive_text() 4 ctx1efgh

69 assert data == "foo" 4 ctx1efgh

70 data = websocket.receive_text() 4 ctx1efgh

71 assert data == "bar" 4 ctx1efgh

72 data = websocket.receive_text() 4 ctx1efgh

73 assert data == "baz" 4 ctx1efgh

74 

75 

76def test_websocket_dependency_after_yield_broken(): 4 ctx1abcd

77 with pytest.raises(ValueError, match="Session closed"): 4 ctx1iljk

78 with client.websocket_connect("/ws-broken"): 4 ctx1iljk

79 pass # pragma no cover 3 ctx1ijk