Coverage for farmbot_sidecar_starter_pack/functions/broker.py: 100%

191 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-09-11 15:43 -0700

1""" 

2BrokerConnect class. 

3""" 

4 

5# └── functions/broker.py 

6# ├── [BROKER] connect() 

7# ├── [BROKER] disconnect() 

8# ├── [BROKER] publish() 

9# ├── [BROKER] start_listen() 

10# ├── [BROKER] stop_listen() 

11# └── [BROKER] listen() 

12 

13import time 

14import math 

15import json 

16import uuid 

17from datetime import datetime 

18import paho.mqtt.client as mqtt 

19 

20 

21class BrokerConnect(): 

22 """Broker connection class.""" 

23 

24 def __init__(self, state): 

25 self.state = state 

26 self.client = None 

27 

28 def connect(self): 

29 """Establish persistent connection to send messages via message broker.""" 

30 

31 self.state.check_token() 

32 

33 self.client = mqtt.Client() 

34 self.client.username_pw_set( 

35 username=self.state.token["token"]["unencoded"]["bot"], 

36 password=self.state.token["token"]["encoded"] 

37 ) 

38 

39 self.client.connect( 

40 self.state.token["token"]["unencoded"]["mqtt"], 

41 port=1883, 

42 keepalive=60 

43 ) 

44 

45 self.client.loop_start() 

46 

47 self.state.print_status(description="Connected to message broker.") 

48 

49 def disconnect(self): 

50 """Disconnect from the message broker.""" 

51 

52 if self.client is not None: 

53 self.client.loop_stop() 

54 self.client.disconnect() 

55 description = "Disconnected from message broker." 

56 self.state.print_status(description=description) 

57 

58 def wrap_message(self, message, priority=None): 

59 """Wrap message in CeleryScript format.""" 

60 rpc = { 

61 "kind": "rpc_request", 

62 "args": { 

63 "label": "", 

64 }, 

65 "body": [message], 

66 } 

67 

68 if priority is not None: 

69 rpc["args"]["priority"] = priority 

70 

71 return rpc 

72 

73 def publish(self, message): 

74 """Publish messages containing CeleryScript via the message broker.""" 

75 

76 if self.client is None: 

77 self.connect() 

78 

79 rpc = message 

80 if rpc["kind"] != "rpc_request": 

81 rpc = self.wrap_message(rpc) 

82 

83 if rpc["args"]["label"] == "": 

84 if self.state.test_env: 

85 rpc["args"]["label"] = "test" 

86 else: 

87 rpc["args"]["label"] = uuid.uuid4().hex 

88 

89 self.state.print_status(description="Publishing to 'from_clients'") 

90 self.state.print_status(endpoint_json=rpc, update_only=True) 

91 if self.state.dry_run: 

92 self.state.print_status( 

93 description="Sending disabled, message not sent.", 

94 update_only=True) 

95 else: 

96 self.listen("from_device", publish_payload=rpc) 

97 

98 response = self.state.last_messages.get("from_device", []) 

99 if len(response) > 0: 

100 if response[-1]["kind"] == "rpc_ok": 

101 self.state.print_status( 

102 description="Success response received.", 

103 update_only=True) 

104 self.state.error = None 

105 else: 

106 self.state.print_status( 

107 description="Error response received.", 

108 update_only=True) 

109 self.state.error = "RPC error response received." 

110 

111 self.state.last_published = rpc 

112 

113 def start_listen(self, channel="#", message_options=None): 

114 """Establish persistent subscription to message broker channels.""" 

115 options = message_options or {} 

116 path = (options.get("path", "") or "").split(".") or [] 

117 path = [key for key in path if key != ""] 

118 diff_only = options.get("diff_only") 

119 

120 if self.client is None: 

121 self.connect() 

122 

123 def add_message(key, value): 

124 """Add message to last_messages.""" 

125 if key not in self.state.last_messages: 

126 self.state.last_messages[key] = [] 

127 self.state.last_messages[key].append(value) 

128 

129 # Set on_message callback 

130 def on_message(_client, _userdata, msg): 

131 """on_message callback""" 

132 channel_key = msg.topic.split("/")[-1] 

133 payload = json.loads(msg.payload) 

134 

135 if channel == "#": 

136 add_message(channel, payload) 

137 add_message(channel_key, payload) 

138 

139 for key in path: 

140 payload = payload[key] 

141 path_channel = f"{channel_key}_excerpt" 

142 add_message(path_channel, payload) 

143 

144 if diff_only: 

145 diff = payload 

146 key = path_channel if len(path) > 0 else channel_key 

147 last_messages = self.state.last_messages.get(key, []) 

148 if len(last_messages) > 1: 

149 current = last_messages[-1] 

150 previous = last_messages[-2] 

151 diff, _is_different = difference(current, previous) 

152 payload = diff 

153 add_message(f"{channel_key}_diffs", payload) 

154 

155 self.state.print_status(description="", update_only=True) 

156 description = "New message" 

157 if len(path) > 0: 

158 description += f" {'.'.join(path)}" 

159 if diff_only: 

160 description += " diff" 

161 description += f" from {msg.topic}" 

162 timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S") 

163 description += f" ({timestamp})" 

164 self.state.print_status( 

165 endpoint_json=payload, 

166 description=description) 

167 

168 self.client.on_message = on_message 

169 

170 # Subscribe to channel 

171 device_id_str = self.state.token["token"]["unencoded"]["bot"] 

172 self.client.subscribe(f"bot/{device_id_str}/{channel}") 

173 description = f"Connected to message broker channel '{channel}'" 

174 if channel == "#": 

175 description = "Connected to all message broker channels" 

176 self.state.print_status(description=description) 

177 

178 # Start listening 

179 self.client.loop_start() 

180 description = f"Now listening to message broker channel '{channel}'" 

181 if channel == "#": 

182 description = "Now listening to all message broker channels" 

183 self.state.print_status(description=description) 

184 

185 def stop_listen(self): 

186 """End subscription to all message broker channels.""" 

187 

188 self.client.loop_stop() 

189 

190 self.state.print_status( 

191 description="Stopped listening to all message broker channels.") 

192 

193 def listen(self, 

194 channel="#", 

195 duration=None, 

196 publish_payload=None, 

197 stop_count=1, 

198 message_options=None): 

199 """Listen to a message broker channel for the provided duration in seconds.""" 

200 publish = publish_payload is not None 

201 message = (publish_payload or {}).get("body", [{}])[0] 

202 # Prepare duration option 

203 timeout_key = "listen" 

204 if message.get("kind") in ["move", "find_home", "calibrate"]: 

205 timeout_key = "movements" 

206 duration_seconds = duration or self.state.timeout[timeout_key] 

207 if message.get("kind") == "wait": 

208 duration_seconds += message["args"]["milliseconds"] / 1000 

209 if stop_count > 1: 

210 duration_seconds = math.inf 

211 # Prepare label option 

212 label = None 

213 if publish and publish_payload["args"]["label"] != "": 

214 label = publish_payload["args"]["label"] 

215 if message.get("kind") == "read_status": 

216 # Getting the RPC response to read_status isn't as important as 

217 # returning the status as soon as possible, since the device 

218 # will still publish a status within a few seconds even if the 

219 # read_status command isn't received. 

220 channel = "status" 

221 label = None 

222 

223 # Print status message 

224 description = "Listening to message broker" 

225 if channel != "#": 

226 description += f" channel '{channel}'" 

227 if duration_seconds != math.inf: 

228 description += f" for {duration_seconds} seconds" 

229 if label is not None: 

230 description += f" for label '{label}'" 

231 plural = "s are" if stop_count > 1 else " is" 

232 description += f" until {stop_count} message{plural} received" 

233 description += "..." 

234 self.state.print_status(description=description) 

235 

236 # Start listening 

237 start_time = datetime.now() 

238 self.start_listen(channel, message_options) 

239 if not self.state.test_env: 

240 if channel == "#": 

241 self.state.last_messages = {"#": []} 

242 else: 

243 self.state.last_messages[channel] = [] 

244 if publish: 

245 time.sleep(0.1) # wait for start_listen to be ready 

246 device_id_str = self.state.token["token"]["unencoded"]["bot"] 

247 publish_topic = f"bot/{device_id_str}/from_clients" 

248 self.client.publish( 

249 publish_topic, 

250 payload=json.dumps(publish_payload)) 

251 self.state.print_status(update_only=True, description="", end="") 

252 while (datetime.now() - start_time).seconds < duration_seconds: 

253 self.state.print_status(update_only=True, description=".", end="") 

254 time.sleep(0.25) 

255 last_messages = self.state.last_messages.get(channel, []) 

256 if len(last_messages) > 0: 

257 # If a label is provided, verify the label matches 

258 if label is not None and last_messages[-1]["args"]["label"] != label: 

259 self.state.last_messages[channel] = [] 

260 continue 

261 if len(last_messages) > (stop_count - 1): 

262 seconds = (datetime.now() - start_time).seconds 

263 prefix = f"{stop_count} messages" 

264 if stop_count == 1: 

265 prefix = "Message" 

266 description = f"{prefix} received after {seconds} seconds" 

267 self.state.print_status( 

268 description=description, 

269 update_only=True) 

270 break 

271 if len(self.state.last_messages.get(channel, [])) == 0: 

272 self.state.print_status(description="", update_only=True) 

273 secs = duration_seconds 

274 description = f"Did not receive message after {secs} seconds" 

275 self.state.print_status( 

276 description=description, 

277 update_only=True) 

278 self.state.error = "Timed out waiting for RPC response." 

279 

280 self.stop_listen() 

281 

282 

283def difference(next_state, prev_state): 

284 """Find the difference between two states.""" 

285 is_different = False 

286 diff = {} 

287 

288 for key, next_value in next_state.items(): 

289 if key not in prev_state: 

290 diff[key] = next_value 

291 is_different = True 

292 continue 

293 prev_value = prev_state[key] 

294 if next_value != prev_value: 

295 if isinstance(next_value, dict) and isinstance(prev_value, dict): 

296 nested_diff, nested_is_different = difference( 

297 next_value, 

298 prev_value) 

299 if nested_is_different: 

300 diff[key] = nested_diff 

301 is_different = True 

302 else: 

303 diff[key] = next_value 

304 is_different = True 

305 

306 for key in prev_state: 

307 if key not in next_state: 

308 is_different = True 

309 

310 return diff, is_different