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
« prev ^ index » next coverage.py v7.4.4, created at 2024-09-04 17:51 -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
20class BrokerConnect():
21 """Broker connection class."""
22 def __init__(self, state):
23 self.state = state
24 self.client = None
26 def connect(self):
27 """Establish persistent connection to send messages via message broker."""
29 self.state.check_token()
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 )
37 self.client.connect(
38 self.state.token["token"]["unencoded"]["mqtt"],
39 port=1883,
40 keepalive=60
41 )
43 self.client.loop_start()
45 self.state.print_status(description="Connected to message broker.")
47 def disconnect(self):
48 """Disconnect from the message broker."""
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.")
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 }
65 if priority is not None:
66 rpc["args"]["priority"] = priority
68 return rpc
70 def publish(self, message):
71 """Publish messages containing CeleryScript via the message broker."""
73 if self.client is None:
74 self.connect()
76 rpc = message
77 if rpc["kind"] != "rpc_request":
78 rpc = self.wrap_message(rpc)
80 if rpc["args"]["label"] == "":
81 rpc["args"]["label"] = uuid.uuid4().hex if not self.state.test_env else "test"
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)
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."
99 self.state.last_published = rpc
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")
108 if self.client is None:
109 self.connect()
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)
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)
123 if channel == "#":
124 add_message(channel, payload)
125 add_message(channel_key, payload)
127 for key in path:
128 payload = payload[key]
129 path_channel = f"{channel_key}_excerpt"
130 add_message(path_channel, payload)
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)
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)
156 self.client.on_message = on_message
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)
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)
173 def stop_listen(self):
174 """End subscription to all message broker channels."""
176 self.client.loop_stop()
178 self.state.print_status(description="Stopped listening to all message broker channels.")
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
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)
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."
255 self.stop_listen()
257def difference(next_state, prev_state):
258 """Find the difference between two states."""
259 is_different = False
260 diff = {}
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
278 for key in prev_state:
279 if key not in next_state:
280 is_different = True
282 return diff, is_different