Coverage for backend/tests/test_ws.py: 98%

1import pytest 

2from fastapi.testclient import TestClient 

3 

4from backend.dependencies import LobbyManager, lobby_manager 

5from backend.hello import app 

6 

7lm = LobbyManager(["a", "b", "c"]) 

8 

9 

10def lobby_manager_override() -> LobbyManager: 

11 """Override lobby manager.""" 

12 return lm 

13 

14 

15app.dependency_overrides[lobby_manager] = lobby_manager_override 

16 

17 

18@pytest.fixture 

19def lobby_client() -> TestClient: 

20 """Create a TestClient.""" 

21 client = TestClient(app) 

22 return client 

23 

24 

25@pytest.mark.slow 

26def test_websocket(lobby_client: TestClient) -> None: 

27 """Test that websocket connection does not hang.""" 

28 code = lobby_client.post("lobby").json()["code"] 

29 player_1_id = lobby_client.post("/lobby/join", json={"code": code}).json()["id"] 

30 player_2_id = lobby_client.post("/lobby/join", json={"code": code}).json()["id"] 

31 

32 with lobby_client.websocket_connect("/ws") as websocket_1, lobby_client.websocket_connect("/ws") as websocket_2: 

33 websocket_1.send_json(dict(data=dict(type="initializer", code=code, id=player_1_id))) 

34 websocket_2.send_json(dict(data=dict(type="initializer", code=code, id=player_2_id))) 

35 

36 # since player_1 joins first, this should be a player count message with count=1 

37 player_count = websocket_1.receive_json() 

38 assert player_count["data"]["count"] == 1 

39 # the next message for player_1 will be the player count update for player_2 joining 

40 player_count = websocket_1.receive_json() 

41 assert player_count["data"]["count"] == 2 

42 

43 # whereas player_2 should only get the player count for player_2 joining 

44 player_count = websocket_2.receive_json() 

45 assert player_count["data"]["count"] == 2 

46 

47 # player_1 starts the game... 

48 websocket_1.send_json(dict(data=dict(type="lobby_lifecycle", lifecycle_type="game_start"))) 

49 

50 # ...so player_2 receives the start broadcast 

51 data = websocket_2.receive_json()["data"] 

52 assert data["lifecycle_type"] == "game_start" 

53 

54 # next message is the role assignment 

55 player_1_role = websocket_1.receive_json()["data"] 

56 _ = websocket_2.receive_json()["data"] 

57 

58 if player_1_role["role"] == "manager": 

59 manager = websocket_1 

60 cook = websocket_2 

61 else: 

62 manager = websocket_2 

63 cook = websocket_1 

64 

65 # next message to the manager should be the order 

66 data = manager.receive_json()["data"] 

67 order = data["order"] 

68 burger = order["burger"]["ingredients"] 

69 assert burger[0] == "Bottom Bun" 

70 assert burger[-1] == "Top Bun" 

71 

72 # cook should send an order component... 

73 cook.send_json( 

74 dict( 

75 data=dict( 

76 type="game_state", 

77 game_state_update_type="order_component", 

78 component_type="burger", 

79 component=dict(ingredients=burger), 

80 ) 

81 ) 

82 ) 

83 

84 # ...which should be sent to the manager 

85 data = manager.receive_json()["data"] 

86 assert data["component"]["ingredients"] == burger 

87 

88 # manager submits order 

89 manager.send_json( 

90 dict( 

91 data=dict( 

92 type="game_state", 

93 game_state_update_type="order_submission", 

94 order=dict(burger=dict(ingredients=burger), side=None, drink=None), 

95 ) 

96 ) 

97 ) 

98 

99 # and gets a score 

100 data = manager.receive_json()["data"] 

101 

102 # a perfect burger gets 500 

103 assert data["score"] == 500 

104 

105 

106@pytest.mark.slow 

107def test_websocket_full_order(lobby_client: TestClient) -> None: 

108 """Test that websocket handles all 4 players and full order flow.""" 

109 code = lobby_client.post("lobby").json()["code"] 

110 

111 player_ids = [lobby_client.post("/lobby/join", json={"code": code}).json()["id"] for _ in range(4)] 

112 websockets = [lobby_client.websocket_connect("/ws") for _ in range(4)] 

113 

114 with websockets[0] as ws1, websockets[1] as ws2, websockets[2] as ws3, websockets[3] as ws4: 

115 sockets = [ws1, ws2, ws3, ws4] 

116 

117 # Send initializer messages 

118 for ws, pid in zip(sockets, player_ids, strict=False): 

119 ws.send_json({"data": {"type": "initializer", "code": code, "id": pid}}) 

120 

121 # Each websocket receives player count updates 

122 for ws in sockets: 

123 while True: 

124 msg = ws.receive_json() 

125 if msg["data"].get("count") == 4: 

126 break 

127 

128 # Start game from first player 

129 ws1.send_json({"data": {"type": "lobby_lifecycle", "lifecycle_type": "game_start"}}) 

130 

131 # Wait for game start broadcast and role assignments 

132 roles = {} 

133 for ws, _ in zip(sockets, player_ids, strict=False): 

134 data = ws.receive_json()["data"] 

135 if data.get("lifecycle_type") == "game_start": 

136 data = ws.receive_json()["data"] 

137 roles[data["role"]] = ws 

138 

139 assert {"manager", "burger", "side", "drink"} == set(roles) 

140 

141 manager = roles["manager"] 

142 burger_cook = roles["burger"] 

143 side_cook = roles["side"] 

144 drink_cook = roles["drink"] 

145 

146 # Manager receives an order 

147 data = manager.receive_json()["data"] 

148 order = data["order"] 

149 burger = order["burger"]["ingredients"] 

150 side = order["side"]["table_state"] 

151 drink = order["drink"] 

152 

153 # Burger cook sends component 

154 burger_cook.send_json( 

155 { 

156 "data": { 

157 "type": "game_state", 

158 "game_state_update_type": "order_component", 

159 "component_type": "burger", 

160 "component": {"ingredients": burger}, 

161 } 

162 } 

163 ) 

164 data = manager.receive_json()["data"] 

165 assert data["component"]["ingredients"] == burger 

166 

167 # Side cook sends component 

168 side_cook.send_json( 

169 { 

170 "data": { 

171 "type": "game_state", 

172 "game_state_update_type": "order_component", 

173 "component_type": "side", 

174 "component": {"table_state": side}, 

175 } 

176 } 

177 ) 

178 data = manager.receive_json()["data"] 

179 assert data["component"]["table_state"] == side 

180 

181 # Drink cook sends component 

182 drink_cook.send_json( 

183 { 

184 "data": { 

185 "type": "game_state", 

186 "game_state_update_type": "order_component", 

187 "component_type": "drink", 

188 "component": {"color": drink["color"], "fill": drink["fill"], "size": drink["size"]}, 

189 } 

190 } 

191 ) 

192 data = manager.receive_json()["data"] 

193 assert data["component"]["color"] == drink["color"] 

194 

195 # Manager submits complete order 

196 manager.send_json( 

197 { 

198 "data": { 

199 "type": "game_state", 

200 "game_state_update_type": "order_submission", 

201 "order": { 

202 "burger": {"ingredients": burger}, 

203 "side": {"table_state": side}, 

204 "drink": {"color": drink["color"], "fill": drink["fill"], "size": drink["size"]}, 

205 }, 

206 } 

207 } 

208 ) 

209 

210 # Manager receives score 

211 score_msg = manager.receive_json()["data"] 

212 assert "score" in score_msg 

213 assert score_msg["score"] == 1200 

214 

215 

216@pytest.mark.slow 

217def test_websocket_spam_chat(lobby_client: TestClient) -> None: 

218 """Test that a bunch of chat messages are received and broadcast without pausing.""" 

219 code = lobby_client.post("lobby").json()["code"] 

220 player_1_id = lobby_client.post("/lobby/join", json={"code": code}).json()["id"] 

221 player_2_id = lobby_client.post("/lobby/join", json={"code": code}).json()["id"] 

222 

223 with lobby_client.websocket_connect("/ws") as websocket_1, lobby_client.websocket_connect("/ws") as websocket_2: 

224 websocket_1.send_json(dict(data=dict(type="initializer", code=code, id=player_1_id))) 

225 websocket_2.send_json(dict(data=dict(type="initializer", code=code, id=player_2_id))) 

226 

227 # since player_1 joins first, this should be a player count message with count=1 

228 player_count = websocket_1.receive_json() 

229 assert player_count["data"]["count"] == 1 

230 # the next message for player_1 will be the player count update for player_2 joining 

231 player_count = websocket_1.receive_json() 

232 assert player_count["data"]["count"] == 2 

233 

234 # whereas player_2 should only get the player count for player_2 joining 

235 player_count = websocket_2.receive_json() 

236 assert player_count["data"]["count"] == 2 

237 

238 for _ in range(10): 

239 websocket_1.send_json(dict(data=dict(type="chat", typing=True, id=player_1_id))) 

240 

241 for _ in range(10): 

242 data = websocket_2.receive_json()["data"] 

243 assert data["type"] == "chat"