save last messages by channel

Author: gabrielburnworth

functions/broker.py

@@ -87,10 +87,10 @@ def on_connect(self, _client, _userdata, _flags, _rc, channel):
 
         self.state.print_status(description=f"Connected to message broker channel {channel}")
 
-    def on_message(self, _client, _userdata, msg):
+    def on_message(self, _client, _userdata, msg, channel):
         """Callback function when message received from message broker."""
 
-        self.state.last_message = json.loads(msg.payload)
+        self.state.last_messages[channel] = json.loads(msg.payload)
 
         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")
 
@@ -100,9 +100,16 @@ def start_listen(self, channel="#"):
         if self.client is None:
             self.connect()
 
-        # Wrap on_connect to pass channel argument
-        self.client.on_connect = lambda client, userdata, flags, rc: self.on_connect(client, userdata, flags, rc, channel)
-        self.client.on_message = self.on_message
+        def on_connect(client, userdata, flags, rc):
+            """Wrap on_connect to pass channel argument."""
+            self.on_connect(client, userdata, flags, rc, channel)
+
+        def on_message(client, userdata, msg):
+            """Wrap on_message to pass channel argument."""
+            self.on_message(client, userdata, msg, channel)
+
+        self.client.on_connect = on_connect
+        self.client.on_message = on_message
 
         self.client.loop_start()
         self.state.print_status(description=f"Now listening to message broker channel {channel}.")
@@ -121,17 +128,17 @@ def listen(self, duration, channel):
         start_time = datetime.now()
         self.start_listen(channel)
         if not self.state.test_env:
-            self.state.last_message = None
+            self.state.last_messages[channel] = None
         while (datetime.now() - start_time).seconds < duration:
             self.state.print_status(update_only=True, description=".", end="")
             time.sleep(0.25)
-            if self.state.last_message is not None:
+            if self.state.last_messages.get(channel) is not None:
                 seconds = (datetime.now() - start_time).seconds
                 self.state.print_status(
                     description=f"Message received after {seconds} seconds",
                     update_only=True)
                 break
-        if self.state.last_message is None:
+        if self.state.last_messages.get(channel) is None:
             self.state.print_status(
                 description=f"Did not receive message after {duration} seconds",
                 update_only=True)

functions/information.py

@@ -137,7 +137,7 @@ def read_status(self):
 
         self.broker.listen(self.state.broker_listen_duration, "status")
 
-        status_tree = self.state.last_message
+        status_tree = self.state.last_messages.get("status")
 
         self.state.print_status(update_only=True, endpoint_json=status_tree)
         return status_tree

state.py

@@ -24,7 +24,7 @@ class State():
     def __init__(self):
         self.token = None
         self.error = None
-        self.last_message = None
+        self.last_messages = {}
         self.verbosity = 2
         self.broker_listen_duration = 15
         self.test_env = False

tests/tests_main.py

@@ -581,10 +581,10 @@ def test_listen_clear_last(self, mock_mqtt):
         '''Test listen command: clear last message'''
         mock_client = Mock()
         mock_mqtt.return_value = mock_client
-        self.fb.state.last_message = "message"
+        self.fb.state.last_messages = {'#': "message"}
         self.fb.state.test_env = False
         self.fb.listen(1)
-        self.assertIsNone(self.fb.state.last_message)
+        self.assertIsNone(self.fb.state.last_messages['#'])
 
     @patch('requests.request')
     @patch('paho.mqtt.client.Client')
@@ -1074,7 +1074,7 @@ def exec_command():
     def test_get_xyz(self):
         '''Test get_xyz command'''
         def exec_command():
-            self.fb.state.last_message = {
+            self.fb.state.last_messages['status'] = {
                 'location_data': {'position': {'x': 1, 'y': 2, 'z': 3}},
             }
             position = self.fb.get_xyz()
@@ -1091,7 +1091,7 @@ def exec_command():
     def test_get_xyz_no_status(self):
         '''Test get_xyz command: no status'''
         def exec_command():
-            self.fb.state.last_message = None
+            self.fb.state.last_messages['status'] = None
             position = self.fb.get_xyz()
             self.assertIsNone(position)
         self.send_command_test_helper(
@@ -1106,7 +1106,7 @@ def exec_command():
     def test_check_position(self):
         '''Test check_position command: at position'''
         def exec_command():
-            self.fb.state.last_message = {
+            self.fb.state.last_messages['status'] = {
                 'location_data': {'position': {'x': 1, 'y': 2, 'z': 3}},
             }
             at_position = self.fb.check_position(1, 2, 3, 0)
@@ -1123,7 +1123,7 @@ def exec_command():
     def test_check_position_false(self):
         '''Test check_position command: not at position'''
         def exec_command():
-            self.fb.state.last_message = {
+            self.fb.state.last_messages['status'] = {
                 'location_data': {'position': {'x': 1, 'y': 2, 'z': 3}},
             }
             at_position = self.fb.check_position(0, 0, 0, 2)
@@ -1140,7 +1140,7 @@ def exec_command():
     def test_check_position_no_status(self):
         '''Test check_position command: no status'''
         def exec_command():
-            self.fb.state.last_message = None
+            self.fb.state.last_messages['status'] = None
             at_position = self.fb.check_position(0, 0, 0, 2)
             self.assertFalse(at_position)
         self.send_command_test_helper(
@@ -1344,7 +1344,7 @@ def test_get_seed_tray_cell_invalid_pullout_direction(self):
     def test_get_job_one(self):
         '''Test get_job command: get one job'''
         def exec_command():
-            self.fb.state.last_message = {
+            self.fb.state.last_messages['status'] = {
                 'jobs': {
                     'job name': {'status': 'working'},
                 },
@@ -1363,7 +1363,7 @@ def exec_command():
     def test_get_job_all(self):
         '''Test get_job command: get all jobs'''
         def exec_command():
-            self.fb.state.last_message = {
+            self.fb.state.last_messages['status'] = {
                 'jobs': {
                     'job name': {'status': 'working'},
                 },
@@ -1382,7 +1382,7 @@ def exec_command():
     def test_get_job_no_status(self):
         '''Test get_job command: no status'''
         def exec_command():
-            self.fb.state.last_message = None
+            self.fb.state.last_messages['status'] = None
             job = self.fb.get_job('job name')
             self.assertIsNone(job)
         self.send_command_test_helper(