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
« prev ^ index » next coverage.py v7.4.4, created at 2024-09-11 15:43 -0700
1"""
2BrokerConnect class.
3"""
5# └── functions/broker.py
6# ├── [BROKER] connect()
7# ├── [BROKER] disconnect()
8# ├── [BROKER] publish()
9# ├── [BROKER] start_listen()
10# ├── [BROKER] stop_listen()
11# └── [BROKER] listen()
13import time
14import math
15import json
16import uuid
17from datetime import datetime
18import paho.mqtt.client as mqtt
21class BrokerConnect():
22 """Broker connection class."""
24 def __init__(self, state):
25 self.state = state
26 self.client = None
28 def connect(self):
29 """Establish persistent connection to send messages via message broker."""
31 self.state.check_token()
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 )
39 self.client.connect(
40 self.state.token["token"]["unencoded"]["mqtt"],
41 port=1883,
42 keepalive=60
43 )
45 self.client.loop_start()
47 self.state.print_status(description="Connected to message broker.")
49 def disconnect(self):
50 """Disconnect from the message broker."""
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)
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 }
68 if priority is not None:
69 rpc["args"]["priority"] = priority
71 return rpc
73 def publish(self, message):
74 """Publish messages containing CeleryScript via the message broker."""
76 if self.client is None:
77 self.connect()
79 rpc = message
80 if rpc["kind"] != "rpc_request":
81 rpc = self.wrap_message(rpc)
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
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)
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."
111 self.state.last_published = rpc
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")
120 if self.client is None:
121 self.connect()
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)
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)
135 if channel == "#":
136 add_message(channel, payload)
137 add_message(channel_key, payload)
139 for key in path:
140 payload = payload[key]
141 path_channel = f"{channel_key}_excerpt"
142 add_message(path_channel, payload)
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)
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)
168 self.client.on_message = on_message
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)
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)
185 def stop_listen(self):
186 """End subscription to all message broker channels."""
188 self.client.loop_stop()
190 self.state.print_status(
191 description="Stopped listening to all message broker channels.")
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
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)
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."
280 self.stop_listen()
283def difference(next_state, prev_state):
284 """Find the difference between two states."""
285 is_different = False
286 diff = {}
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
306 for key in prev_state:
307 if key not in next_state:
308 is_different = True
310 return diff, is_different