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 18 ctx1qhkrensiltfoujmvgp

15 

16 

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

18 s = Session() 18 ctx1qhkrensiltfoujmvgp

19 yield s 18 ctx1qhkrensiltfoujmvgp

20 s.open = False 18 ctx1qhkrensiltfoujmvgp

21 global_state = global_context.get() 18 ctx1qhkrensiltfoujmvgp

22 global_state["session_closed"] = True 18 ctx1qhkrensiltfoujmvgp

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 9 ctx1hkeilfjmg

33 self.open = True 9 ctx1hkeilfjmg

34 

35 

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

37 assert session is session_b 3 ctx1efg

38 named_session = NamedSession(name="named") 3 ctx1efg

39 yield named_session, session_b 3 ctx1efg

40 named_session.open = False 3 ctx1efg

41 global_state = global_context.get() 3 ctx1efg

42 global_state["named_session_closed"] = True 3 ctx1efg

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") 3 ctx1hij

50 yield named_session, session 3 ctx1hij

51 named_session.open = False 3 ctx1hij

52 global_state = global_context.get() 3 ctx1hij

53 global_state["named_func_session_closed"] = True 3 ctx1hij

54 

55 

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

57 named_session = NamedSession(name="named") 3 ctx1klm

58 return named_session, session 3 ctx1klm

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() 3 ctx1qsu

78 await websocket.send_json({"is_open": session.open}) 3 ctx1qsu

79 

80 

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

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

83 await websocket.accept() 3 ctx1rtv

84 await websocket.send_json({"is_open": session.open}) 3 ctx1rtv

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() 3 ctx1nop

94 await websocket.send_json( 3 ctx1nop

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() 3 ctx1efg

102 await websocket.send_json( 3 ctx1efg

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() 3 ctx1hij

112 await websocket.send_json( 3 ctx1hij

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() 3 ctx1klm

122 await websocket.send_json( 3 ctx1klm

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({}) 3 ctx1qsu

132 global_state = global_context.get() 3 ctx1qsu

133 with client.websocket_connect("/function-scope") as websocket: 3 ctx1qsu

134 data = websocket.receive_json() 3 ctx1qsu

135 assert data["is_open"] is True 3 ctx1qsu

136 assert global_state["session_closed"] is True 3 ctx1qsu

137 

138 

139def test_request_scope() -> None: 4 ctx1abcd

140 global_context.set({}) 3 ctx1rtv

141 global_state = global_context.get() 3 ctx1rtv

142 with client.websocket_connect("/request-scope") as websocket: 3 ctx1rtv

143 data = websocket.receive_json() 3 ctx1rtv

144 assert data["is_open"] is True 3 ctx1rtv

145 assert global_state["session_closed"] is True 3 ctx1rtv

146 

147 

148def test_two_scopes() -> None: 4 ctx1abcd

149 global_context.set({}) 3 ctx1nop

150 global_state = global_context.get() 3 ctx1nop

151 with client.websocket_connect("/two-scopes") as websocket: 3 ctx1nop

152 data = websocket.receive_json() 3 ctx1nop

153 assert data["func_is_open"] is True 3 ctx1nop

154 assert data["req_is_open"] is True 3 ctx1nop

155 assert global_state["session_closed"] is True 3 ctx1nop

156 

157 

158def test_sub() -> None: 4 ctx1abcd

159 global_context.set({}) 3 ctx1efg

160 global_state = global_context.get() 3 ctx1efg

161 with client.websocket_connect("/sub") as websocket: 3 ctx1efg

162 data = websocket.receive_json() 3 ctx1efg

163 assert data["named_session_open"] is True 3 ctx1efg

164 assert data["session_open"] is True 3 ctx1efg

165 assert global_state["session_closed"] is True 3 ctx1efg

166 assert global_state["named_session_closed"] is True 3 ctx1efg

167 

168 

169def test_broken_scope() -> None: 4 ctx1abcd

170 with pytest.raises( 3 ctx1wxy

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") 3 ctx1wxy

176 async def get_broken( 3 ctx1wxy

177 websocket: WebSocket, sessions: BrokenSessionsDep 3 ctx1wxy

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

179 pass 

180 

181 

182def test_named_function_scope() -> None: 4 ctx1abcd

183 global_context.set({}) 3 ctx1hij

184 global_state = global_context.get() 3 ctx1hij

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

186 data = websocket.receive_json() 3 ctx1hij

187 assert data["named_session_open"] is True 3 ctx1hij

188 assert data["session_open"] is True 3 ctx1hij

189 assert global_state["session_closed"] is True 3 ctx1hij

190 assert global_state["named_func_session_closed"] is True 3 ctx1hij

191 

192 

193def test_regular_function_scope() -> None: 4 ctx1abcd

194 global_context.set({}) 3 ctx1klm

195 global_state = global_context.get() 3 ctx1klm

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

197 data = websocket.receive_json() 3 ctx1klm

198 assert data["named_session_open"] is True 3 ctx1klm

199 assert data["session_open"] is True 3 ctx1klm

200 assert global_state["session_closed"] is True 3 ctx1klm