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

183 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-09-04 17:51 -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 

20class BrokerConnect(): 

21 """Broker connection class.""" 

22 def __init__(self, state): 

23 self.state = state 

24 self.client = None 

25 

26 def connect(self): 

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

28 

29 self.state.check_token() 

30 

31 self.client = mqtt.Client() 

32 self.client.username_pw_set( 

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

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

35 ) 

36 

37 self.client.connect( 

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

39 port=1883, 

40 keepalive=60 

41 ) 

42 

43 self.client.loop_start() 

44 

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

46 

47 def disconnect(self): 

48 """Disconnect from the message broker.""" 

49 

50 if self.client is not None: 

51 self.client.loop_stop() 

52 self.client.disconnect() 

53 self.state.print_status(description="Disconnected from message broker.") 

54 

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

56 """Wrap message in CeleryScript format.""" 

57 rpc = { 

58 "kind": "rpc_request", 

59 "args": { 

60 "label": "", 

61 }, 

62 "body": [message], 

63 } 

64 

65 if priority is not None: 

66 rpc["args"]["priority"] = priority 

67 

68 return rpc 

69 

70 def publish(self, message): 

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

72 

73 if self.client is None: 

74 self.connect() 

75 

76 rpc = message 

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

78 rpc = self.wrap_message(rpc) 

79 

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

81 rpc["args"]["label"] = uuid.uuid4().hex if not self.state.test_env else "test" 

82 

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

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

85 if self.state.dry_run: 

86 self.state.print_status(description="Sending disabled, message not sent.", update_only=True) 

87 else: 

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

89 

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

91 if len(response) > 0: 

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

93 self.state.print_status(description="Success response received.", update_only=True) 

94 self.state.error = None 

95 else: 

96 self.state.print_status(description="Error response received.", update_only=True) 

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

98 

99 self.state.last_published = rpc 

100 

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

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

103 options = message_options or {} 

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

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

106 diff_only = options.get("diff_only") 

107 

108 if self.client is None: 

109 self.connect() 

110 

111 def add_message(key, value): 

112 """Add message to last_messages.""" 

113 if key not in self.state.last_messages: 

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

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

116 

117 # Set on_message callback 

118 def on_message(_client, _userdata, msg): 

119 """on_message callback""" 

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

121 payload = json.loads(msg.payload) 

122 

123 if channel == "#": 

124 add_message(channel, payload) 

125 add_message(channel_key, payload) 

126 

127 for key in path: 

128 payload = payload[key] 

129 path_channel = f"{channel_key}_excerpt" 

130 add_message(path_channel, payload) 

131 

132 if diff_only: 

133 diff = payload 

134 prev_channel_key = path_channel if len(path) > 0 else channel_key 

135 last_messages = self.state.last_messages.get(prev_channel_key, []) 

136 if len(last_messages) > 1: 

137 current = last_messages[-1] 

138 previous = last_messages[-2] 

139 diff, _is_different = difference(current, previous) 

140 payload = diff 

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

142 

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

144 description = "New message" 

145 if len(path) > 0: 

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

147 if diff_only: 

148 description += " diff" 

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

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

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

152 self.state.print_status( 

153 endpoint_json=payload, 

154 description=description) 

155 

156 self.client.on_message = on_message 

157 

158 # Subscribe to channel 

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

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

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

162 if channel == "#": 

163 description = "Connected to all message broker channels" 

164 self.state.print_status(description=description) 

165 

166 # Start listening 

167 self.client.loop_start() 

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

169 if channel == "#": 

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

171 self.state.print_status(description=description) 

172 

173 def stop_listen(self): 

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

175 

176 self.client.loop_stop() 

177 

178 self.state.print_status(description="Stopped listening to all message broker channels.") 

179 

180 def listen(self, channel="#", duration=None, publish_payload=None, stop_count=1, message_options=None): 

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

182 publish = publish_payload is not None 

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

184 # Prepare duration option 

185 timeout_key = "listen" 

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

187 timeout_key = "movements" 

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

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

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

191 if stop_count > 1: 

192 duration_seconds = math.inf 

193 # Prepare label option 

194 label = None 

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

196 label = publish_payload["args"]["label"] 

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

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

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

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

201 # read_status command isn't received. 

202 channel = "status" 

203 label = None 

204 

205 # Print status message 

206 description = "Listening to message broker" 

207 if channel != "#": 

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

209 if duration_seconds != math.inf: 

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

211 if label is not None: 

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

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

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

215 description += "..." 

216 self.state.print_status(description=description) 

217 

218 # Start listening 

219 start_time = datetime.now() 

220 self.start_listen(channel, message_options) 

221 if not self.state.test_env: 

222 if channel == "#": 

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

224 else: 

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

226 if publish: 

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

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

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

230 self.client.publish(publish_topic, payload=json.dumps(publish_payload)) 

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

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

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

234 time.sleep(0.25) 

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

236 if len(last_messages) > 0: 

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

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

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

240 continue 

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

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

243 prefix = "Message" if stop_count == 1 else f"{stop_count} messages" 

244 self.state.print_status( 

245 description=f"{prefix} received after {seconds} seconds", 

246 update_only=True) 

247 break 

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

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

250 self.state.print_status( 

251 description=f"Did not receive message after {duration_seconds} seconds", 

252 update_only=True) 

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

254 

255 self.stop_listen() 

256 

257def difference(next_state, prev_state): 

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

259 is_different = False 

260 diff = {} 

261 

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

263 if key not in prev_state: 

264 diff[key] = next_value 

265 is_different = True 

266 continue 

267 prev_value = prev_state[key] 

268 if next_value != prev_value: 

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

270 nested_diff, nested_is_different = difference(next_value, prev_value) 

271 if nested_is_different: 

272 diff[key] = nested_diff 

273 is_different = True 

274 else: 

275 diff[key] = next_value 

276 is_different = True 

277 

278 for key in prev_state: 

279 if key not in next_state: 

280 is_different = True 

281 

282 return diff, is_different