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

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 json 

15import uuid 

16from datetime import datetime 

17import paho.mqtt.client as mqtt 

18 

19class BrokerConnect(): 

20 """Broker connection class.""" 

21 def __init__(self, state): 

22 self.state = state 

23 self.client = None 

24 

25 def connect(self): 

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

27 

28 self.state.check_token() 

29 

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 ) 

35 

36 self.client.connect( 

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

38 port=1883, 

39 keepalive=60 

40 ) 

41 

42 self.client.loop_start() 

43 

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

45 

46 def disconnect(self): 

47 """Disconnect from the message broker.""" 

48 

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.") 

53 

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 } 

63 

64 if priority is not None: 

65 rpc["args"]["priority"] = priority 

66 

67 return rpc 

68 

69 def publish(self, message): 

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

71 

72 if self.client is None: 

73 self.connect() 

74 

75 rpc = message 

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

77 rpc = self.wrap_message(rpc) 

78 

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

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

81 

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) 

88 

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." 

97 

98 self.state.last_published = rpc 

99 

100 def start_listen(self, channel="#"): 

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

102 

103 if self.client is None: 

104 self.connect() 

105 

106 # Set on_message callback 

107 def on_message(_client, _userdata, msg): 

108 """on_message callback""" 

109 

110 self.state.last_messages[channel] = json.loads(msg.payload) 

111 

112 

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})") 

118 

119 self.client.on_message = on_message 

120 

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}'") 

125 

126 # Start listening 

127 self.client.loop_start() 

128 self.state.print_status(description=f"Now listening to message broker channel '{channel}'.") 

129 

130 def stop_listen(self): 

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

132 

133 self.client.loop_stop() 

134 

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

136 

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 

158 

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) 

165 

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." 

197 

198 self.stop_listen()