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

105 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-08-31 14:00 -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"TOPIC: {msg.topic} ({timestamp})\n") 

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 duration_seconds = duration or self.state.broker_listen_duration 

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

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

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

144 publish = publish_payload is not None 

145 label = None 

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

147 label = publish_payload["args"]["label"] 

148 

149 # Print status message 

150 channel_str = f" channel '{channel}'" if channel != "#" else "" 

151 duration_str = f" for {duration_seconds} seconds" 

152 label_str = f" for label '{label}'" if label is not None else "" 

153 description = f"Listening to message broker{channel_str}{duration_str}{label_str}..." 

154 self.state.print_status(description=description) 

155 

156 # Start listening 

157 start_time = datetime.now() 

158 self.start_listen(channel) 

159 if not self.state.test_env: 

160 self.state.last_messages[channel] = None 

161 if publish: 

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

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

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

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

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

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

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

169 time.sleep(0.25) 

170 last_message = self.state.last_messages.get(channel) 

171 if last_message is not None: 

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

173 if label is not None and last_message["args"]["label"] != label: 

174 self.state.last_messages[channel] = None 

175 continue 

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

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

178 self.state.print_status( 

179 description=f"Message received after {seconds} seconds", 

180 update_only=True) 

181 break 

182 if self.state.last_messages.get(channel) is None: 

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

184 self.state.print_status( 

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

186 update_only=True) 

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

188 

189 self.stop_listen()