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

6from fastapi.responses import StreamingResponse 7 ctx1abcdefg

7from fastapi.testclient import TestClient 7 ctx1abcdefg

8from typing_extensions import Annotated 7 ctx1abcdefg

9 

10 

11class Session: 7 ctx1abcdefg

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

13 self.data = ["foo", "bar", "baz"] 49 ctx1JohpKCXLqirMDYNsjtOEZPukvQF0RwlxSG1TymzUH2VAnBWI3

14 self.open = True 49 ctx1JohpKCXLqirMDYNsjtOEZPukvQF0RwlxSG1TymzUH2VAnBWI3

15 

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

17 for item in self.data: 42 ctx1JohpKCLqirMDNsjtOEPukvQFRwlxSGTymzUHVAnBWI

18 if self.open: 42 ctx1JohpKCLqirMDNsjtOEPukvQFRwlxSGTymzUHVAnBWI

19 yield item 14 ctx1KCMDOEQFSGUHWI

20 else: 

21 raise ValueError("Session closed") 28 ctx1JohpLqirNsjtPukvRwlxTymzVAnB

22 

23 

24@contextmanager 7 ctx1abcdefg

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

26 session = Session() 49 ctx1JohpKCXLqirMDYNsjtOEZPukvQF0RwlxSG1TymzUH2VAnBWI3

27 try: 49 ctx1JohpKCXLqirMDYNsjtOEZPukvQF0RwlxSG1TymzUH2VAnBWI3

28 yield session 49 ctx1JohpKCXLqirMDYNsjtOEZPukvQF0RwlxSG1TymzUH2VAnBWI3

29 finally: 

30 session.open = False 49 ctx1JohpKCXLqirMDYNsjtOEZPukvQF0RwlxSG1TymzUH2VAnBWI3

31 

32 

33def dep_session() -> Any: 7 ctx1abcdefg

34 with acquire_session() as s: 21 ctx1KCXMDYOEZQF0SG1UH2WI3

35 yield s 21 ctx1KCXMDYOEZQF0SG1UH2WI3

36 

37 

38def broken_dep_session() -> Any: 7 ctx1abcdefg

39 with acquire_session() as s: 28 ctx1JohpLqirNsjtPukvRwlxTymzVAnB

40 s.open = False 28 ctx1JohpLqirNsjtPukvRwlxTymzVAnB

41 yield s 28 ctx1JohpLqirNsjtPukvRwlxTymzVAnB

42 

43 

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

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

46 

47app = FastAPI() 7 ctx1abcdefg

48 

49 

50@app.get("/data") 7 ctx1abcdefg

51def get_data(session: SessionDep) -> Any: 7 ctx1abcdefg

52 data = list(session) 7 ctx1KMOQSUW

53 return data 7 ctx1KMOQSUW

54 

55 

56@app.get("/stream-simple") 7 ctx1abcdefg

57def get_stream_simple(session: SessionDep) -> Any: 7 ctx1abcdefg

58 def iter_data(): 7 ctx1XYZ0123

59 yield from ["x", "y", "z"] 7 ctx1XYZ0123

60 

61 return StreamingResponse(iter_data()) 7 ctx1XYZ0123

62 

63 

64@app.get("/stream-session") 7 ctx1abcdefg

65def get_stream_session(session: SessionDep) -> Any: 7 ctx1abcdefg

66 def iter_data(): 7 ctx1CDEFGHI

67 yield from session 7 ctx1CDEFGHI

68 

69 return StreamingResponse(iter_data()) 7 ctx1CDEFGHI

70 

71 

72@app.get("/broken-session-data") 7 ctx1abcdefg

73def get_broken_session_data(session: BrokenSessionDep) -> Any: 7 ctx1abcdefg

74 return list(session) 14 ctx1JoLqNsPuRwTyVA

75 

76 

77@app.get("/broken-session-stream") 7 ctx1abcdefg

78def get_broken_session_stream(session: BrokenSessionDep) -> Any: 7 ctx1abcdefg

79 def iter_data(): 14 ctx1hpirjtkvlxmznB

80 yield from session 14 ctx1hpirjtkvlxmznB

81 

82 return StreamingResponse(iter_data()) 14 ctx1hpirjtkvlxmznB

83 

84 

85client = TestClient(app) 7 ctx1abcdefg

86 

87 

88def test_regular_no_stream(): 7 ctx1abcdefg

89 response = client.get("/data") 7 ctx1KMOQSUW

90 assert response.json() == ["foo", "bar", "baz"] 7 ctx1KMOQSUW

91 

92 

93def test_stream_simple(): 7 ctx1abcdefg

94 response = client.get("/stream-simple") 7 ctx1XYZ0123

95 assert response.text == "xyz" 7 ctx1XYZ0123

96 

97 

98def test_stream_session(): 7 ctx1abcdefg

99 response = client.get("/stream-session") 7 ctx1CDEFGHI

100 assert response.text == "foobarbaz" 7 ctx1CDEFGHI

101 

102 

103def test_broken_session_data(): 7 ctx1abcdefg

104 with pytest.raises(ValueError, match="Session closed"): 7 ctx1JLNPRTV

105 client.get("/broken-session-data") 7 ctx1JLNPRTV

106 

107 

108def test_broken_session_data_no_raise(): 7 ctx1abcdefg

109 client = TestClient(app, raise_server_exceptions=False) 7 ctx1oqsuwyA

110 response = client.get("/broken-session-data") 7 ctx1oqsuwyA

111 assert response.status_code == 500 7 ctx1oqsuwyA

112 assert response.text == "Internal Server Error" 7 ctx1oqsuwyA

113 

114 

115def test_broken_session_stream_raise(): 7 ctx1abcdefg

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

117 with pytest.raises((ValueError, Exception)): 7 ctx1prtvxzB

118 client.get("/broken-session-stream") 7 ctx1prtvxzB

119 

120 

121def test_broken_session_stream_no_raise(): 7 ctx1abcdefg

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: 7 ctx1hijklmn

128 response = client.get("/broken-session-stream") 7 ctx1hijklmn

129 assert response.status_code == 200 7 ctx1hijklmn

130 assert response.text == "" 7 ctx1hijklmn