Replace mqtt tuple with dataclass

Author: andchiind

src/isar/models/events.py

@@ -14,7 +14,7 @@
 from robot_interface.models.mission.mission import Mission
 from robot_interface.models.mission.status import RobotStatus
 from robot_interface.models.mission.task import InspectionTask
-from robot_interface.telemetry.mqtt_client import MQTTQueueType
+from robot_interface.telemetry.mqtt_client import MQTTQueueMessage
 
 
 class EmptyMessage:
@@ -92,7 +92,7 @@ def __init__(self) -> None:
             "uploader", maxsize=10
         )
 
-        self.mqtt_queue: Queue[MQTTQueueType] = Queue[MQTTQueueType]()
+        self.mqtt_queue: Queue[MQTTQueueMessage] = Queue[MQTTQueueMessage]()
 
         self.state: Event[States] = Event("state")
 

src/isar/services/service_connections/mqtt/mqtt_client.py

@@ -14,7 +14,7 @@
 from paho.mqtt.reasoncodes import ReasonCode
 
 from isar.config.settings import settings
-from robot_interface.telemetry.mqtt_client import MqttClientInterface
+from robot_interface.telemetry.mqtt_client import MqttClientInterface, MQTTQueueMessage
 
 
 def props_expiry(seconds: int) -> Properties:
@@ -44,10 +44,10 @@ def _on_giveup(data: Details) -> None:
 
 
 class MqttClient(MqttClientInterface):
-    def __init__(self, mqtt_queue: Queue) -> None:
+    def __init__(self, mqtt_queue: Queue[MQTTQueueMessage]) -> None:
         self.logger = logging.getLogger("mqtt_client")
         self.logger.setLevel("INFO")
-        self.mqtt_queue: Queue = mqtt_queue
+        self.mqtt_queue: Queue[MQTTQueueMessage] = mqtt_queue
 
         username: str = settings.MQTT_USERNAME
         password: str = ""
@@ -93,23 +93,16 @@ def run(self) -> None:
                 time.sleep(0)  # avoid CPU spin
                 continue
             try:
-                item: tuple[str, str, int, bool, Properties | None] = (
-                    self.mqtt_queue.get(timeout=1)
-                )
-                if len(item) == 4:
-                    topic, payload, qos, retain = item
-                    properties = None
-                else:
-                    topic, payload, qos, retain, properties = item
+                item: MQTTQueueMessage = self.mqtt_queue.get(timeout=1)
             except Empty:
                 continue
 
             self.publish(
-                topic=topic,
-                payload=payload,
-                qos=qos,
-                retain=retain,
-                properties=properties,
+                topic=item.topic,
+                payload=item.payload,
+                qos=item.qos,
+                retain=item.retain,
+                properties=item.properties,
             )
 
     def on_connect(

src/isar/services/service_connections/mqtt/robot_heartbeat_publisher.py

@@ -4,12 +4,12 @@
 
 from isar.config.settings import settings
 from isar.services.service_connections.mqtt.mqtt_client import props_expiry
-from robot_interface.telemetry.mqtt_client import MqttPublisher
+from robot_interface.telemetry.mqtt_client import MqttPublisher, MQTTQueueMessage
 from robot_interface.telemetry.payloads import RobotHeartbeatPayload
 
 
 class RobotHeartbeatPublisher:
-    def __init__(self, mqtt_queue: Queue):
+    def __init__(self, mqtt_queue: Queue[MQTTQueueMessage]):
         self.mqtt_publisher: MqttPublisher = MqttPublisher(mqtt_queue=mqtt_queue)
 
     def run(self) -> None:

src/isar/services/service_connections/mqtt/robot_info_publisher.py

@@ -3,12 +3,12 @@
 from queue import Queue
 
 from isar.config.settings import robot_settings, settings
-from robot_interface.telemetry.mqtt_client import MqttPublisher
+from robot_interface.telemetry.mqtt_client import MqttPublisher, MQTTQueueMessage
 from robot_interface.telemetry.payloads import RobotInfoPayload
 
 
 class RobotInfoPublisher:
-    def __init__(self, mqtt_queue: Queue):
+    def __init__(self, mqtt_queue: Queue[MQTTQueueMessage]):
         self.mqtt_publisher: MqttPublisher = MqttPublisher(mqtt_queue=mqtt_queue)
 
     def run(self) -> None:

src/isar/services/utilities/mqtt_utilities.py

@@ -10,7 +10,7 @@
 from robot_interface.telemetry.mqtt_client import (
     MqttClientInterface,
     MqttPublisher,
-    MQTTQueueType,
+    MQTTQueueMessage,
 )
 from robot_interface.telemetry.payloads import (
     InterventionNeededPayload,
@@ -50,7 +50,7 @@ def publish_task_status(
 
 
 def publish_mission_status(
-    mqtt_queue: Queue[MQTTQueueType],
+    mqtt_queue: Queue[MQTTQueueMessage],
     mission_id: str,
     mission_status: MissionStatus,
     error_message: ErrorMessage | None,
@@ -97,7 +97,7 @@ def publish_isar_status(
 
 
 def publish_mission_aborted(
-    mqtt_queue: Queue[MQTTQueueType], current_mission_id: str | None, reason: str
+    mqtt_queue: Queue[MQTTQueueMessage], current_mission_id: str | None, reason: str
 ) -> None:
     mqtt_publisher: MqttPublisher = MqttPublisher(mqtt_queue=mqtt_queue)
 
@@ -119,7 +119,7 @@ def publish_mission_aborted(
 
 
 def publish_intervention_needed(
-    mqtt_queue: Queue[MQTTQueueType], error_message: str
+    mqtt_queue: Queue[MQTTQueueMessage], error_message: str
 ) -> None:
     """Publishes the intervention needed message to the MQTT Broker"""
     mqtt_publisher: MqttPublisher = MqttPublisher(mqtt_queue=mqtt_queue)

src/robot_interface/robot_interface.py

@@ -8,7 +8,10 @@
 from robot_interface.models.mission.status import MissionStatus, RobotStatus, TaskStatus
 from robot_interface.models.mission.task import InspectionTask
 from robot_interface.models.robots.media import MediaConfig
-from robot_interface.telemetry.mqtt_client import MqttTelemetryPublisher
+from robot_interface.telemetry.mqtt_client import (
+    MQTTQueueMessage,
+    MqttTelemetryPublisher,
+)
 
 
 class RobotInterface(metaclass=ABCMeta):
@@ -216,7 +219,7 @@ def generate_media_config(self) -> MediaConfig | None:
 
     @abstractmethod
     def get_telemetry_publishers(
-        self, queue: Queue, isar_id: str, robot_name: str
+        self, queue: Queue[MQTTQueueMessage], isar_id: str, robot_name: str
     ) -> list[MqttTelemetryPublisher]:
         """
         Set up telemetry publisher threads to publish regular updates for pose, battery

src/robot_interface/telemetry/mqtt_client.py

@@ -3,6 +3,7 @@
 import time
 from abc import ABCMeta, abstractmethod
 from collections.abc import Callable
+from dataclasses import dataclass
 from datetime import UTC, datetime
 from logging import Logger
 from queue import Queue
@@ -19,7 +20,14 @@
 )
 from robot_interface.telemetry.payloads import CloudHealthPayload
 
-MQTTQueueType = tuple[str, str, int, bool, Properties | None]
+
+@dataclass
+class MQTTQueueMessage:
+    topic: str
+    payload: str
+    qos: int
+    retain: bool
+    properties: Properties | None
 
 
 def props_expiry(seconds: int) -> Properties:
@@ -56,8 +64,8 @@ def publish(
 
 
 class MqttPublisher(MqttClientInterface):
-    def __init__(self, mqtt_queue: Queue[MQTTQueueType]) -> None:
-        self.mqtt_queue: Queue[MQTTQueueType] = mqtt_queue
+    def __init__(self, mqtt_queue: Queue[MQTTQueueMessage]) -> None:
+        self.mqtt_queue: Queue[MQTTQueueMessage] = mqtt_queue
 
     def publish(
         self,
@@ -67,12 +75,12 @@ def publish(
         retain: bool = False,
         properties: Properties | None = None,
     ) -> None:
-        queue_message: tuple[str, str, int, bool, Properties | None] = (
-            topic,
-            payload,
-            qos,
-            retain,
-            properties,
+        queue_message: MQTTQueueMessage = MQTTQueueMessage(
+            topic=topic,
+            payload=payload,
+            qos=qos,
+            retain=retain,
+            properties=properties,
         )
         self.mqtt_queue.put(queue_message)
 
@@ -81,13 +89,13 @@ class MqttTelemetryPublisher(Thread):
     def __init__(
         self,
         name: str,
-        mqtt_queue: Queue[MQTTQueueType],
+        mqtt_queue: Queue[MQTTQueueMessage],
         telemetry_method: Callable,
         topic: str,
         interval: float,
         should_expire: bool,
     ) -> None:
-        self.mqtt_queue: Queue[MQTTQueueType] = mqtt_queue
+        self.mqtt_queue: Queue[MQTTQueueMessage] = mqtt_queue
         self.telemetry_method: Callable = telemetry_method
         self.topic: str = topic
         self.interval: float = interval
@@ -130,11 +138,11 @@ def run(self) -> None:
             if self.should_expire:
                 properties = props_expiry(settings.MQTT_TELEMETRY_EXPIRY)
 
-            queue_message: MQTTQueueType = (
-                topic,
-                payload,
-                0,
-                False,
-                properties,
+            queue_message: MQTTQueueMessage = MQTTQueueMessage(
+                topic=topic,
+                payload=payload,
+                qos=0,
+                retain=False,
+                properties=properties,
             )
             self.mqtt_queue.put(queue_message)

tests/isar/state_machine/states/test_going_to_lockdown_state.py

@@ -60,7 +60,7 @@ def test_stopping_lockdown_transitions_to_going_to_lockdown(events: Events) -> N
     assert not events.mqtt_queue.empty()
     mqtt_message = events.mqtt_queue.get(block=False)
     assert mqtt_message is not None
-    mqtt_payload_topic = mqtt_message[0]
+    mqtt_payload_topic = mqtt_message.topic
     assert mqtt_payload_topic is settings.TOPIC_ISAR_MISSION_ABORTED
 
 

tests/test_mocks/robot_interface.py

@@ -17,7 +17,10 @@
 from robot_interface.models.mission.task import InspectionTask
 from robot_interface.models.robots.media import MediaConfig, MediaConnectionType
 from robot_interface.robot_interface import RobotInterface
-from robot_interface.telemetry.mqtt_client import MqttTelemetryPublisher
+from robot_interface.telemetry.mqtt_client import (
+    MQTTQueueMessage,
+    MqttTelemetryPublisher,
+)
 from tests.test_mocks.inspection import stub_image_metadata
 
 _ROBOT_FRAME = Frame(name="robot")
@@ -85,7 +88,7 @@ def register_inspection_callback(
         return
 
     def get_telemetry_publishers(
-        self, queue: Queue, isar_id: str, robot_name: str
+        self, queue: Queue[MQTTQueueMessage], isar_id: str, robot_name: str
     ) -> list[MqttTelemetryPublisher]:
         return []