Coverage for functions/broker.py: 100%

75 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-08-29 17:51 -0700

1""" 

2BrokerConnect class. 

3""" 

4 

5# └── functions/broker.py 

6# ├── [BROKER] connect() 

7# ├── [BROKER] disconnect() 

8# ├── [BROKER] publish() 

9# ├── [BROKER] on_connect() 

10# ├── [BROKER] on_message 

11# ├── [BROKER] start_listen() 

12# ├── [BROKER] stop_listen() 

13# └── [BROKER] listen() 

14 

15import time 

16import json 

17from datetime import datetime 

18import paho.mqtt.client as mqtt 

19 

20class BrokerConnect(): 

21 """Broker connection class.""" 

22 def __init__(self, state): 

23 self.state = state 

24 self.client = None 

25 

26 def connect(self): 

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

28 

29 self.state.check_token() 

30 

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 ) 

36 

37 self.client.connect( 

38 self.state.token['token']['unencoded']['mqtt'], 

39 port=1883, 

40 keepalive=60 

41 ) 

42 

43 self.client.loop_start() 

44 

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

46 

47 def disconnect(self): 

48 """Disconnect from the message broker.""" 

49 

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

54 

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 } 

64 

65 if priority is not None: 

66 rpc['args']['priority'] = priority 

67 

68 return rpc 

69 

70 def publish(self, message): 

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

72 

73 if self.client is None: 

74 self.connect() 

75 

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

77 message = self.wrap_message(message) 

78 

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

80 topic = f"bot/{device_id_str}/from_clients" 

81 if not self.state.dry_run: 

82 self.client.publish(topic, payload=json.dumps(message)) 

83 self.state.print_status(description=f"Publishing to {topic}:") 

84 self.state.print_status(endpoint_json=message, update_only=True) 

85 if self.state.dry_run: 

86 self.state.print_status(description="Sending disabled, message not sent.", update_only=True) 

87 

88 def on_connect(self, _client, _userdata, _flags, _rc, channel): 

89 """Callback function when connection to message broker is successful.""" 

90 

91 self.client.subscribe( 

92 f"bot/{self.state.token['token']['unencoded']['bot']}/{channel}") 

93 

94 self.state.print_status(description=f"Connected to message broker channel {channel}") 

95 

96 def on_message(self, _client, _userdata, msg, channel): 

97 """Callback function when message received from message broker.""" 

98 

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

100 

101 self.state.print_status(endpoint_json=json.loads(msg.payload), description=f"TOPIC: {msg.topic} ({datetime.now().strftime('%Y-%m-%d %H:%M:%S')})\n") 

102 

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

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

105 

106 if self.client is None: 

107 self.connect() 

108 

109 def on_connect(client, userdata, flags, rc): 

110 """Wrap on_connect to pass channel argument.""" 

111 self.on_connect(client, userdata, flags, rc, channel) 

112 

113 def on_message(client, userdata, msg): 

114 """Wrap on_message to pass channel argument.""" 

115 self.on_message(client, userdata, msg, channel) 

116 

117 self.client.on_connect = on_connect 

118 self.client.on_message = on_message 

119 

120 self.client.loop_start() 

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

122 

123 def stop_listen(self): 

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

125 

126 self.client.loop_stop() 

127 self.client.disconnect() 

128 

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

130 

131 def listen(self, duration, channel): 

132 """Listen to a message broker channel for the provided duration in seconds.""" 

133 self.state.print_status(description=f"Listening to message broker for {duration} seconds...") 

134 start_time = datetime.now() 

135 self.start_listen(channel) 

136 if not self.state.test_env: 

137 self.state.last_messages[channel] = None 

138 while (datetime.now() - start_time).seconds < duration: 

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

140 time.sleep(0.25) 

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

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

143 self.state.print_status( 

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

145 update_only=True) 

146 break 

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

148 self.state.print_status( 

149 description=f"Did not receive message after {duration} seconds", 

150 update_only=True) 

151 

152 self.stop_listen()