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

6from fastapi.responses import StreamingResponse 6 ctx1abcdef

7from fastapi.testclient import TestClient 6 ctx1abcdef

8from typing_extensions import Annotated 6 ctx1abcdef

9 

10 

11class Session: 6 ctx1abcdef

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

13 self.data = ["foo", "bar", "baz"] 42 ctx1EmgnFyQGohpHzRIqirJASKsjtLBTMukvNCUOwlxPDV

14 self.open = True 42 ctx1EmgnFyQGohpHzRIqirJASKsjtLBTMukvNCUOwlxPDV

15 

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

17 for item in self.data: 36 ctx1EmgnFyGohpHzIqirJAKsjtLBMukvNCOwlxPD

18 if self.open: 36 ctx1EmgnFyGohpHzIqirJAKsjtLBMukvNCOwlxPD

19 yield item 12 ctx1FyHzJALBNCPD

20 else: 

21 raise ValueError("Session closed") 24 ctx1EmgnGohpIqirKsjtMukvOwlx

22 

23 

24@contextmanager 6 ctx1abcdef

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

26 session = Session() 42 ctx1EmgnFyQGohpHzRIqirJASKsjtLBTMukvNCUOwlxPDV

27 try: 42 ctx1EmgnFyQGohpHzRIqirJASKsjtLBTMukvNCUOwlxPDV

28 yield session 42 ctx1EmgnFyQGohpHzRIqirJASKsjtLBTMukvNCUOwlxPDV

29 finally: 

30 session.open = False 42 ctx1EmgnFyQGohpHzRIqirJASKsjtLBTMukvNCUOwlxPDV

31 

32 

33def dep_session() -> Any: 6 ctx1abcdef

34 with acquire_session() as s: 18 ctx1FyQHzRJASLBTNCUPDV

35 yield s 18 ctx1FyQHzRJASLBTNCUPDV

36 

37 

38def broken_dep_session() -> Any: 6 ctx1abcdef

39 with acquire_session() as s: 24 ctx1EmgnGohpIqirKsjtMukvOwlx

40 s.open = False 24 ctx1EmgnGohpIqirKsjtMukvOwlx

41 yield s 24 ctx1EmgnGohpIqirKsjtMukvOwlx

42 

43 

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

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

46 

47app = FastAPI() 6 ctx1abcdef

48 

49 

50@app.get("/data") 6 ctx1abcdef

51def get_data(session: SessionDep) -> Any: 6 ctx1abcdef

52 data = list(session) 6 ctx1FHJLNP

53 return data 6 ctx1FHJLNP

54 

55 

56@app.get("/stream-simple") 6 ctx1abcdef

57def get_stream_simple(session: SessionDep) -> Any: 6 ctx1abcdef

58 def iter_data(): 6 ctx1QRSTUV

59 yield from ["x", "y", "z"] 6 ctx1QRSTUV

60 

61 return StreamingResponse(iter_data()) 6 ctx1QRSTUV

62 

63 

64@app.get("/stream-session") 6 ctx1abcdef

65def get_stream_session(session: SessionDep) -> Any: 6 ctx1abcdef

66 def iter_data(): 6 ctx1yzABCD

67 yield from session 6 ctx1yzABCD

68 

69 return StreamingResponse(iter_data()) 6 ctx1yzABCD

70 

71 

72@app.get("/broken-session-data") 6 ctx1abcdef

73def get_broken_session_data(session: BrokenSessionDep) -> Any: 6 ctx1abcdef

74 return list(session) 12 ctx1EmGoIqKsMuOw

75 

76 

77@app.get("/broken-session-stream") 6 ctx1abcdef

78def get_broken_session_stream(session: BrokenSessionDep) -> Any: 6 ctx1abcdef

79 def iter_data(): 12 ctx1gnhpirjtkvlx

80 yield from session 12 ctx1gnhpirjtkvlx

81 

82 return StreamingResponse(iter_data()) 12 ctx1gnhpirjtkvlx

83 

84 

85client = TestClient(app) 6 ctx1abcdef

86 

87 

88def test_regular_no_stream(): 6 ctx1abcdef

89 response = client.get("/data") 6 ctx1FHJLNP

90 assert response.json() == ["foo", "bar", "baz"] 6 ctx1FHJLNP

91 

92 

93def test_stream_simple(): 6 ctx1abcdef

94 response = client.get("/stream-simple") 6 ctx1QRSTUV

95 assert response.text == "xyz" 6 ctx1QRSTUV

96 

97 

98def test_stream_session(): 6 ctx1abcdef

99 response = client.get("/stream-session") 6 ctx1yzABCD

100 assert response.text == "foobarbaz" 6 ctx1yzABCD

101 

102 

103def test_broken_session_data(): 6 ctx1abcdef

104 with pytest.raises(ValueError, match="Session closed"): 6 ctx1EGIKMO

105 client.get("/broken-session-data") 6 ctx1EGIKMO

106 

107 

108def test_broken_session_data_no_raise(): 6 ctx1abcdef

109 client = TestClient(app, raise_server_exceptions=False) 6 ctx1moqsuw

110 response = client.get("/broken-session-data") 6 ctx1moqsuw

111 assert response.status_code == 500 6 ctx1moqsuw

112 assert response.text == "Internal Server Error" 6 ctx1moqsuw

113 

114 

115def test_broken_session_stream_raise(): 6 ctx1abcdef

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

117 with pytest.raises((ValueError, Exception)): 6 ctx1nprtvx

118 client.get("/broken-session-stream") 6 ctx1nprtvx

119 

120 

121def test_broken_session_stream_no_raise(): 6 ctx1abcdef

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: 6 ctx1ghijkl

128 response = client.get("/broken-session-stream") 6 ctx1ghijkl

129 assert response.status_code == 200 6 ctx1ghijkl

130 assert response.text == "" 6 ctx1ghijkl