Coverage for farmbot_sidecar_starter_pack/functions/broker.py: 100%
110 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-09-01 12:19 -0700
« prev ^ index » next coverage.py v7.4.4, created at 2024-09-01 12:19 -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 json
15import uuid
16from datetime import datetime
17import paho.mqtt.client as mqtt
19class BrokerConnect():
20 """Broker connection class."""
21 def __init__(self, state):
22 self.state = state
23 self.client = None
25 def connect(self):
26 """Establish persistent connection to send messages via message broker."""
28 self.state.check_token()
30 self.client = mqtt.Client()
31 self.client.username_pw_set(
32 username=self.state.token["token"]["unencoded"]["bot"],
33 password=self.state.token["token"]["encoded"]
34 )
36 self.client.connect(
37 self.state.token["token"]["unencoded"]["mqtt"],
38 port=1883,
39 keepalive=60
40 )
42 self.client.loop_start()
44 self.state.print_status(description="Connected to message broker.")
46 def disconnect(self):
47 """Disconnect from the message broker."""
49 if self.client is not None:
50 self.client.loop_stop()
51 self.client.disconnect()
52 self.state.print_status(description="Disconnected from message broker.")
54 def wrap_message(self, message, priority=None):
55 """Wrap message in CeleryScript format."""
56 rpc = {
57 "kind": "rpc_request",
58 "args": {
59 "label": "",
60 },
61 "body": [message],
62 }
64 if priority is not None:
65 rpc["args"]["priority"] = priority
67 return rpc
69 def publish(self, message):
70 """Publish messages containing CeleryScript via the message broker."""
72 if self.client is None:
73 self.connect()
75 rpc = message
76 if rpc["kind"] != "rpc_request":
77 rpc = self.wrap_message(rpc)
79 if rpc["args"]["label"] == "":
80 rpc["args"]["label"] = uuid.uuid4().hex if not self.state.test_env else "test"
82 self.state.print_status(description="Publishing to 'from_clients'")
83 self.state.print_status(endpoint_json=rpc, update_only=True)
84 if self.state.dry_run:
85 self.state.print_status(description="Sending disabled, message not sent.", update_only=True)
86 else:
87 self.listen("from_device", publish_payload=rpc)
89 response = self.state.last_messages.get("from_device")
90 if response is not None:
91 if response["kind"] == "rpc_ok":
92 self.state.print_status(description="Success response received.", update_only=True)
93 self.state.error = None
94 else:
95 self.state.print_status(description="Error response received.", update_only=True)
96 self.state.error = "RPC error response received."
98 self.state.last_published = rpc
100 def start_listen(self, channel="#"):
101 """Establish persistent subscription to message broker channels."""
103 if self.client is None:
104 self.connect()
106 # Set on_message callback
107 def on_message(_client, _userdata, msg):
108 """on_message callback"""
110 self.state.last_messages[channel] = json.loads(msg.payload)
113 self.state.print_status(description="", update_only=True)
114 timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
115 self.state.print_status(
116 endpoint_json=json.loads(msg.payload),
117 description=f"New message from {msg.topic} ({timestamp})")
119 self.client.on_message = on_message
121 # Subscribe to channel
122 device_id_str = self.state.token["token"]["unencoded"]["bot"]
123 self.client.subscribe(f"bot/{device_id_str}/{channel}")
124 self.state.print_status(description=f"Connected to message broker channel '{channel}'")
126 # Start listening
127 self.client.loop_start()
128 self.state.print_status(description=f"Now listening to message broker channel '{channel}'.")
130 def stop_listen(self):
131 """End subscription to all message broker channels."""
133 self.client.loop_stop()
135 self.state.print_status(description="Stopped listening to all message broker channels.")
137 def listen(self, channel, duration=None, publish_payload=None):
138 """Listen to a message broker channel for the provided duration in seconds."""
139 # Prepare parameters
140 message = (publish_payload or {}).get("body", [{}])[0]
141 timeout_key = "listen"
142 if message.get("kind") in ["move", "find_home", "calibrate"]:
143 timeout_key = "movements"
144 duration_seconds = duration or self.state.timeout[timeout_key]
145 if message.get("kind") == "wait":
146 duration_seconds += message["args"]["milliseconds"] / 1000
147 publish = publish_payload is not None
148 label = None
149 if publish and publish_payload["args"]["label"] != "":
150 label = publish_payload["args"]["label"]
151 if message.get("kind") == "read_status":
152 # Getting the RPC response to read_status isn't as important as
153 # returning the status as soon as possible, since the device
154 # will still publish a status within a few seconds even if the
155 # read_status command isn't received.
156 channel = "status"
157 label = None
159 # Print status message
160 channel_str = f" channel '{channel}'" if channel != "#" else ""
161 duration_str = f" for {duration_seconds} seconds"
162 label_str = f" for label '{label}'" if label is not None else ""
163 description = f"Listening to message broker{channel_str}{duration_str}{label_str}..."
164 self.state.print_status(description=description)
166 # Start listening
167 start_time = datetime.now()
168 self.start_listen(channel)
169 if not self.state.test_env:
170 self.state.last_messages[channel] = None
171 if publish:
172 time.sleep(0.1) # wait for start_listen to be ready
173 device_id_str = self.state.token["token"]["unencoded"]["bot"]
174 publish_topic = f"bot/{device_id_str}/from_clients"
175 self.client.publish(publish_topic, payload=json.dumps(publish_payload))
176 self.state.print_status(update_only=True, description="", end="")
177 while (datetime.now() - start_time).seconds < duration_seconds:
178 self.state.print_status(update_only=True, description=".", end="")
179 time.sleep(0.25)
180 last_message = self.state.last_messages.get(channel)
181 if last_message is not None:
182 # If a label is provided, verify the label matches
183 if label is not None and last_message["args"]["label"] != label:
184 self.state.last_messages[channel] = None
185 continue
186 seconds = (datetime.now() - start_time).seconds
187 self.state.print_status(
188 description=f"Message received after {seconds} seconds",
189 update_only=True)
190 break
191 if self.state.last_messages.get(channel) is None:
192 self.state.print_status(description="", update_only=True)
193 self.state.print_status(
194 description=f"Did not receive message after {duration_seconds} seconds",
195 update_only=True)
196 self.state.error = "Timed out waiting for RPC response."
198 self.stop_listen()