Coverage for tests / test_dependency_after_yield_streaming.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 4 ctx1abcd

7from fastapi.responses import StreamingResponse 4 ctx1abcd

8from fastapi.testclient import TestClient 4 ctx1abcd

9 

10 

11class Session: 4 ctx1abcd

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

13 self.data = ["foo", "bar", "baz"] 23 ctx1hqeirnxsyjtfkuozlvgmwpA

14 self.open = True 23 ctx1hqeirnxsyjtfkuozlvgmwpA

15 

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

17 for item in self.data: 20 ctx1hqeirnBsjtfkuolvgmwp

18 if self.open: 20 ctx1hqeirnBsjtfkuolvgmwp

19 yield item 7 ctx1rnsuowp

20 else: 

21 raise ValueError("Session closed") 14 ctx1hqeiCBjtfklvgm

22 

23 

24@contextmanager 4 ctx1abcd

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

26 session = Session() 23 ctx1hqeirnxsyjtfkuozlvgmwpA

27 try: 23 ctx1hqeirnxsyjtfkuozlvgmwpA

28 yield session 23 ctx1hqeirnxsyjtfkuozlvgmwpA

29 finally: 

30 session.open = False 25 ctx1hqeirnxCBsyjtfkuozlvgmwpA

31 

32 

33def dep_session() -> Any: 4 ctx1abcd

34 with acquire_session() as s: 11 ctx1rnxsyuozwpA

35 yield s 11 ctx1rnxsyuozwpA

36 

37 

38def broken_dep_session() -> Any: 4 ctx1abcd

39 with acquire_session() as s: 14 ctx1hqeiCBjtfklvgm

40 s.open = False 14 ctx1hqeiCBjtfklvgm

41 yield s 14 ctx1hqeiCBjtfklvgm

42 

43 

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

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

46 

47app = FastAPI() 4 ctx1abcd

48 

49 

50@app.get("/data") 4 ctx1abcd

51def get_data(session: SessionDep) -> Any: 4 ctx1abcd

52 data = list(session) 4 ctx1rsuw

53 return data 4 ctx1rsuw

54 

55 

56@app.get("/stream-simple") 4 ctx1abcd

57def get_stream_simple(session: SessionDep) -> Any: 4 ctx1abcd

58 def iter_data(): 4 ctx1xyzA

59 yield from ["x", "y", "z"] 4 ctx1xyzA

60 

61 return StreamingResponse(iter_data()) 4 ctx1xyzA

62 

63 

64@app.get("/stream-session") 4 ctx1abcd

65def get_stream_session(session: SessionDep) -> Any: 4 ctx1abcd

66 def iter_data(): 4 ctx1nEop

67 yield from session 4 ctx1nEop

68 

69 return StreamingResponse(iter_data()) 4 ctx1nEop

70 

71 

72@app.get("/broken-session-data") 4 ctx1abcd

73def get_broken_session_data(session: BrokenSessionDep) -> Any: 4 ctx1abcd

74 return list(session) 8 ctx1hqCBjtlv

75 

76 

77@app.get("/broken-session-stream") 4 ctx1abcd

78def get_broken_session_stream(session: BrokenSessionDep) -> Any: 4 ctx1abcd

79 def iter_data(): 8 ctx1eiDFfkgm

80 yield from session 8 ctx1eiDFfkgm

81 

82 return StreamingResponse(iter_data()) 8 ctx1eiDFfkgm

83 

84 

85client = TestClient(app) 4 ctx1abcd

86 

87 

88def test_regular_no_stream(): 4 ctx1abcd

89 response = client.get("/data") 4 ctx1rsuw

90 assert response.json() == ["foo", "bar", "baz"] 4 ctx1rsuw

91 

92 

93def test_stream_simple(): 4 ctx1abcd

94 response = client.get("/stream-simple") 4 ctx1xyzA

95 assert response.text == "xyz" 4 ctx1xyzA

96 

97 

98def test_stream_session(): 4 ctx1abcd

99 response = client.get("/stream-session") 4 ctx1nEop

100 assert response.text == "foobarbaz" 4 ctx1nEop

101 

102 

103def test_broken_session_data(): 4 ctx1abcd

104 with pytest.raises(ValueError, match="Session closed"): 4 ctx1qBtv

105 client.get("/broken-session-data") 4 ctx1qBtv

106 

107 

108def test_broken_session_data_no_raise(): 4 ctx1abcd

109 client = TestClient(app, raise_server_exceptions=False) 4 ctx1hCjl

110 response = client.get("/broken-session-data") 4 ctx1hCjl

111 assert response.status_code == 500 4 ctx1hCjl

112 assert response.text == "Internal Server Error" 4 ctx1hCjl

113 

114 

115def test_broken_session_stream_raise(): 4 ctx1abcd

116 # Can raise ValueError on Pydantic v2 and ExceptionGroup on Pydantic v1 

117 with pytest.raises((ValueError, Exception)): 4 ctx1iFkm

118 client.get("/broken-session-stream") 4 ctx1iFkm

119 

120 

121def test_broken_session_stream_no_raise(): 4 ctx1abcd

122 """ 

123 When a dependency with yield raises after the streaming response already started 

124 the 200 status code is already sent, but there's still an error in the server 

125 afterwards, an exception is raised and captured or shown in the server logs. 

126 """ 

127 with TestClient(app, raise_server_exceptions=False) as client: 4 ctx1eDfg

128 response = client.get("/broken-session-stream") 4 ctx1eDfg

129 assert response.status_code == 200 4 ctx1eDfg

130 assert response.text == "" 4 ctx1eDfg