Coverage for tests / test_dependency_yield_scope_websockets.py: 100%

1from contextvars import ContextVar 4 ctx1abcd

2from typing import Annotated, Any 4 ctx1abcd

3 

4import pytest 4 ctx1abcd

5from fastapi import Depends, FastAPI, WebSocket 4 ctx1abcd

6from fastapi.exceptions import FastAPIError 4 ctx1abcd

7from fastapi.testclient import TestClient 4 ctx1abcd

8 

9global_context: ContextVar[dict[str, Any]] = ContextVar("global_context", default={}) # noqa: B039 4 ctx1abcd

10 

11 

12class Session: 4 ctx1abcd

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

14 self.open = True 20 ctx1shkteouvwilxfpyjmzgq

15 

16 

17async def dep_session() -> Any: 4 ctx1abcd

18 s = Session() 20 ctx1shkteouvwilxfpyjmzgq

19 yield s 20 ctx1shkteouvwilxfpyjmzgq

20 s.open = False 20 ctx1shkteouvwilxfpyjmzgq

21 global_state = global_context.get() 20 ctx1shkteouvwilxfpyjmzgq

22 global_state["session_closed"] = True 20 ctx1shkteouvwilxfpyjmzgq

23 

24 

25SessionFuncDep = Annotated[Session, Depends(dep_session, scope="function")] 4 ctx1abcd

26SessionRequestDep = Annotated[Session, Depends(dep_session, scope="request")] 4 ctx1abcd

27SessionDefaultDep = Annotated[Session, Depends(dep_session)] 4 ctx1abcd

28 

29 

30class NamedSession: 4 ctx1abcd

31 def __init__(self, name: str = "default") -> None: 4 ctx1abcd

32 self.name = name 11 ctx1hkeAnilfjmg

33 self.open = True 11 ctx1hkeAnilfjmg

34 

35 

36def get_named_session(session: SessionRequestDep, session_b: SessionDefaultDep) -> Any: 4 ctx1abcd

37 assert session is session_b 4 ctx1enfg

38 named_session = NamedSession(name="named") 4 ctx1enfg

39 yield named_session, session_b 4 ctx1enfg

40 named_session.open = False 4 ctx1enfg

41 global_state = global_context.get() 4 ctx1enfg

42 global_state["named_session_closed"] = True 4 ctx1enfg

43 

44 

45NamedSessionsDep = Annotated[tuple[NamedSession, Session], Depends(get_named_session)] 4 ctx1abcd

46 

47 

48def get_named_func_session(session: SessionFuncDep) -> Any: 4 ctx1abcd

49 named_session = NamedSession(name="named") 4 ctx1hrij

50 yield named_session, session 4 ctx1hrij

51 named_session.open = False 4 ctx1hrij

52 global_state = global_context.get() 4 ctx1hrij

53 global_state["named_func_session_closed"] = True 4 ctx1hrij

54 

55 

56def get_named_regular_func_session(session: SessionFuncDep) -> Any: 4 ctx1abcd

57 named_session = NamedSession(name="named") 4 ctx1kAlm

58 return named_session, session 4 ctx1kAlm

59 

60 

61BrokenSessionsDep = Annotated[ 4 ctx1abcd

62 tuple[NamedSession, Session], Depends(get_named_func_session) 

63] 

64NamedSessionsFuncDep = Annotated[ 4 ctx1abcd

65 tuple[NamedSession, Session], Depends(get_named_func_session, scope="function") 

66] 

67 

68RegularSessionsDep = Annotated[ 4 ctx1abcd

69 tuple[NamedSession, Session], Depends(get_named_regular_func_session) 

70] 

71 

72app = FastAPI() 4 ctx1abcd

73 

74 

75@app.websocket("/function-scope") 4 ctx1abcd

76async def function_scope(websocket: WebSocket, session: SessionFuncDep) -> Any: 4 ctx1abcd

77 await websocket.accept() 4 ctx1suwy

78 await websocket.send_json({"is_open": session.open}) 4 ctx1suwy

79 

80 

81@app.websocket("/request-scope") 4 ctx1abcd

82async def request_scope(websocket: WebSocket, session: SessionRequestDep) -> Any: 4 ctx1abcd

83 await websocket.accept() 4 ctx1tvxz

84 await websocket.send_json({"is_open": session.open}) 4 ctx1tvxz

85 

86 

87@app.websocket("/two-scopes") 4 ctx1abcd

88async def get_stream_session( 4 ctx1abcd

89 websocket: WebSocket, 

90 function_session: SessionFuncDep, 

91 request_session: SessionRequestDep, 

92) -> Any: 

93 await websocket.accept() 4 ctx1oBpq

94 await websocket.send_json( 4 ctx1oBpq

95 {"func_is_open": function_session.open, "req_is_open": request_session.open} 

96 ) 

97 

98 

99@app.websocket("/sub") 4 ctx1abcd

100async def get_sub(websocket: WebSocket, sessions: NamedSessionsDep) -> Any: 4 ctx1abcd

101 await websocket.accept() 4 ctx1enfg

102 await websocket.send_json( 4 ctx1enfg

103 {"named_session_open": sessions[0].open, "session_open": sessions[1].open} 

104 ) 

105 

106 

107@app.websocket("/named-function-scope") 4 ctx1abcd

108async def get_named_function_scope( 4 ctx1abcd

109 websocket: WebSocket, sessions: NamedSessionsFuncDep 

110) -> Any: 

111 await websocket.accept() 4 ctx1hrij

112 await websocket.send_json( 4 ctx1hrij

113 {"named_session_open": sessions[0].open, "session_open": sessions[1].open} 

114 ) 

115 

116 

117@app.websocket("/regular-function-scope") 4 ctx1abcd

118async def get_regular_function_scope( 4 ctx1abcd

119 websocket: WebSocket, sessions: RegularSessionsDep 

120) -> Any: 

121 await websocket.accept() 4 ctx1kAlm

122 await websocket.send_json( 4 ctx1kAlm

123 {"named_session_open": sessions[0].open, "session_open": sessions[1].open} 

124 ) 

125 

126 

127client = TestClient(app) 4 ctx1abcd

128 

129 

130def test_function_scope() -> None: 4 ctx1abcd

131 global_context.set({}) 4 ctx1suwy

132 global_state = global_context.get() 4 ctx1suwy

133 with client.websocket_connect("/function-scope") as websocket: 4 ctx1suwy

134 data = websocket.receive_json() 4 ctx1suwy

135 assert data["is_open"] is True 4 ctx1suwy

136 assert global_state["session_closed"] is True 4 ctx1suwy

137 

138 

139def test_request_scope() -> None: 4 ctx1abcd

140 global_context.set({}) 4 ctx1tvxz

141 global_state = global_context.get() 4 ctx1tvxz

142 with client.websocket_connect("/request-scope") as websocket: 4 ctx1tvxz

143 data = websocket.receive_json() 4 ctx1tvxz

144 assert data["is_open"] is True 4 ctx1tvxz

145 assert global_state["session_closed"] is True 4 ctx1tvxz

146 

147 

148def test_two_scopes() -> None: 4 ctx1abcd

149 global_context.set({}) 4 ctx1oBpq

150 global_state = global_context.get() 4 ctx1oBpq

151 with client.websocket_connect("/two-scopes") as websocket: 4 ctx1oBpq

152 data = websocket.receive_json() 4 ctx1oBpq

153 assert data["func_is_open"] is True 4 ctx1oBpq

154 assert data["req_is_open"] is True 4 ctx1oBpq

155 assert global_state["session_closed"] is True 4 ctx1oBpq

156 

157 

158def test_sub() -> None: 4 ctx1abcd

159 global_context.set({}) 4 ctx1enfg

160 global_state = global_context.get() 4 ctx1enfg

161 with client.websocket_connect("/sub") as websocket: 4 ctx1enfg

162 data = websocket.receive_json() 4 ctx1enfg

163 assert data["named_session_open"] is True 4 ctx1enfg

164 assert data["session_open"] is True 4 ctx1enfg

165 assert global_state["session_closed"] is True 4 ctx1enfg

166 assert global_state["named_session_closed"] is True 4 ctx1enfg

167 

168 

169def test_broken_scope() -> None: 4 ctx1abcd

170 with pytest.raises( 4 ctx1CFDE

171 FastAPIError, 

172 match='The dependency "get_named_func_session" has a scope of "request", it cannot depend on dependencies with scope "function"', 

173 ): 

174 

175 @app.websocket("/broken-scope") 4 ctx1CFDE

176 async def get_broken( 4 ctx1CFDE

177 websocket: WebSocket, sessions: BrokenSessionsDep 3 ctx1CDE

178 ) -> Any: # pragma: no cover 3 ctx1CDE

179 pass 

180 

181 

182def test_named_function_scope() -> None: 4 ctx1abcd

183 global_context.set({}) 4 ctx1hrij

184 global_state = global_context.get() 4 ctx1hrij

185 with client.websocket_connect("/named-function-scope") as websocket: 4 ctx1hrij

186 data = websocket.receive_json() 4 ctx1hrij

187 assert data["named_session_open"] is True 4 ctx1hrij

188 assert data["session_open"] is True 4 ctx1hrij

189 assert global_state["session_closed"] is True 4 ctx1hrij

190 assert global_state["named_func_session_closed"] is True 4 ctx1hrij

191 

192 

193def test_regular_function_scope() -> None: 4 ctx1abcd

194 global_context.set({}) 4 ctx1kAlm

195 global_state = global_context.get() 4 ctx1kAlm

196 with client.websocket_connect("/regular-function-scope") as websocket: 4 ctx1kAlm

197 data = websocket.receive_json() 4 ctx1kAlm

198 assert data["named_session_open"] is True 4 ctx1kAlm

199 assert data["session_open"] is True 4 ctx1kAlm

200 assert global_state["session_closed"] is True 4 ctx1kAlm