Coverage for functions/broker.py: 100%
75 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-08-29 17:51 -0700
« prev ^ index » next coverage.py v7.4.4, created at 2024-08-29 17:51 -0700
1"""
2BrokerConnect class.
3"""
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()
15import time
16import json
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 if message["kind"] != "rpc_request":
77 message = self.wrap_message(message)
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)
88 def on_connect(self, _client, _userdata, _flags, _rc, channel):
89 """Callback function when connection to message broker is successful."""
91 self.client.subscribe(
92 f"bot/{self.state.token['token']['unencoded']['bot']}/{channel}")
94 self.state.print_status(description=f"Connected to message broker channel {channel}")
96 def on_message(self, _client, _userdata, msg, channel):
97 """Callback function when message received from message broker."""
99 self.state.last_messages[channel] = json.loads(msg.payload)
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")
103 def start_listen(self, channel="#"):
104 """Establish persistent subscription to message broker channels."""
106 if self.client is None:
107 self.connect()
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)
113 def on_message(client, userdata, msg):
114 """Wrap on_message to pass channel argument."""
115 self.on_message(client, userdata, msg, channel)
117 self.client.on_connect = on_connect
118 self.client.on_message = on_message
120 self.client.loop_start()
121 self.state.print_status(description=f"Now listening to message broker channel {channel}.")
123 def stop_listen(self):
124 """End subscription to all message broker channels."""
126 self.client.loop_stop()
127 self.client.disconnect()
129 self.state.print_status(description="Stopped listening to all message broker channels.")
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)
152 self.stop_listen()