Use payloads for mqtt messages

equinor · Dec 20, 2024 · 13e53d3 · 13e53d3
1 parent ce76a64
commit 13e53d3
Show file tree

Hide file tree

Showing 3 changed files with 88 additions and 60 deletions.
diff --git a/src/isar/state_machine/state_machine.py b/src/isar/state_machine/state_machine.py
@@ -34,6 +34,11 @@
 from robot_interface.models.mission.task import TASKS
 from robot_interface.robot_interface import RobotInterface
 from robot_interface.telemetry.mqtt_client import MqttClientInterface
+from robot_interface.telemetry.payloads import (
+    RobotStatusPayload,
+    MissionPayload,
+    TaskPayload,
+)
 from robot_interface.utilities.json_service import EnhancedJSONEncoder
 
 
@@ -482,24 +487,22 @@ def publish_mission_status(self) -> None:
         if self.current_mission:
             if self.current_mission.error_message:
                 error_message = self.current_mission.error_message
-        payload: str = json.dumps(
-            {
-                "isar_id": settings.ISAR_ID,
-                "robot_name": settings.ROBOT_NAME,
-                "mission_id": self.current_mission.id if self.current_mission else None,
-                "status": self.current_mission.status if self.current_mission else None,
-                "error_reason": error_message.error_reason if error_message else None,
-                "error_description": (
-                    error_message.error_description if error_message else None
-                ),
-                "timestamp": datetime.now(timezone.utc),
-            },
-            cls=EnhancedJSONEncoder,
+
+        payload: MissionPayload = MissionPayload(
+            isar_id=settings.ISAR_ID,
+            robot_name=settings.ROBOT_NAME,
+            mission_id=self.current_mission.id if self.current_mission else None,
+            status=self.current_mission.status if self.current_mission else None,
+            error_reason=error_message.error_reason if error_message else None,
+            error_description=(
+                error_message.error_description if error_message else None
+            ),
+            timestamp=datetime.now(timezone.utc),
         )
 
         self.mqtt_publisher.publish(
             topic=settings.TOPIC_ISAR_MISSION,
-            payload=payload,
+            payload=json.dumps(payload, cls=EnhancedJSONEncoder),
             qos=1,
             retain=True,
         )
@@ -514,46 +517,41 @@ def publish_task_status(self, task: TASKS) -> None:
             if task.error_message:
                 error_message = task.error_message
 
-        payload: str = json.dumps(
-            {
-                "isar_id": settings.ISAR_ID,
-                "robot_name": settings.ROBOT_NAME,
-                "mission_id": self.current_mission.id if self.current_mission else None,
-                "task_id": task.id if task else None,
-                "status": task.status if task else None,
-                "task_type": task.type,
-                "error_reason": error_message.error_reason if error_message else None,
-                "error_description": (
-                    error_message.error_description if error_message else None
-                ),
-                "timestamp": datetime.now(timezone.utc),
-            },
-            cls=EnhancedJSONEncoder,
+        payload: TaskPayload = TaskPayload(
+            isar_id=settings.ISAR_ID,
+            robot_name=settings.ROBOT_NAME,
+            mission_id=self.current_mission.id if self.current_mission else None,
+            task_id=task.id if task else None,
+            status=task.status if task else None,
+            task_type=task.type if task else None,
+            error_reason=error_message.error_reason if error_message else None,
+            error_description=(
+                error_message.error_description if error_message else None
+            ),
+            timestamp=datetime.now(timezone.utc),
         )
 
         self.mqtt_publisher.publish(
             topic=settings.TOPIC_ISAR_TASK,
-            payload=payload,
+            payload=json.dumps(payload, cls=EnhancedJSONEncoder),
             qos=1,
             retain=True,
         )
 
     def publish_status(self) -> None:
         if not self.mqtt_publisher:
             return
-        payload: str = json.dumps(
-            {
-                "isar_id": settings.ISAR_ID,
-                "robot_name": settings.ROBOT_NAME,
-                "status": self._current_status(),
-                "timestamp": datetime.now(timezone.utc),
-            },
-            cls=EnhancedJSONEncoder,
+
+        payload: RobotStatusPayload = RobotStatusPayload(
+            isar_id=settings.ISAR_ID,
+            robot_name=settings.ROBOT_NAME,
+            status=self._current_status(),
+            timestamp=datetime.now(timezone.utc),
         )
 
         self.mqtt_publisher.publish(
             topic=settings.TOPIC_ISAR_STATUS,
-            payload=payload,
+            payload=json.dumps(payload, cls=EnhancedJSONEncoder),
             qos=1,
             retain=True,
         )

diff --git a/src/isar/storage/uploader.py b/src/isar/storage/uploader.py
@@ -13,6 +13,7 @@
 from robot_interface.models.inspection.inspection import Inspection
 from robot_interface.models.mission.mission import Mission
 from robot_interface.telemetry.mqtt_client import MqttClientInterface
+from robot_interface.telemetry.payloads import InspectionResultPayload
 from robot_interface.utilities.json_service import EnhancedJSONEncoder
 
 
@@ -149,21 +150,19 @@ def _publish_inspection_result(
         """
         if not self.mqtt_publisher:
             return
-        payload: str = json.dumps(
-            {
-                "isar_id": settings.ISAR_ID,
-                "robot_name": settings.ROBOT_NAME,
-                "inspection_id": inspection.id,
-                "inspection_path": inspection_path,
-                "installation_code": settings.PLANT_SHORT_NAME,
-                "analysis_to_be_run": inspection.metadata.analysis_type,
-                "timestamp": datetime.now(timezone.utc),
-            },
-            cls=EnhancedJSONEncoder,
+
+        payload: InspectionResultPayload = InspectionResultPayload(
+            isar_id=settings.ISAR_ID,
+            robot_name=settings.ROBOT_NAME,
+            inspection_id=inspection.id,
+            inspection_path=inspection_path,
+            installation_code=settings.PLANT_SHORT_NAME,
+            analysis_to_be_run=inspection.metadata.analysis_type,
+            timestamp=datetime.now(timezone.utc),
         )
         self.mqtt_publisher.publish(
             topic=settings.TOPIC_ISAR_INSPECTION_RESULT,
-            payload=payload,
+            payload=json.dumps(payload, cls=EnhancedJSONEncoder),
             qos=1,
             retain=True,
         )
diff --git a/src/robot_interface/telemetry/payloads.py b/src/robot_interface/telemetry/payloads.py
@@ -1,12 +1,13 @@
 from dataclasses import dataclass
 from datetime import datetime
-from typing import List, Optional
+from typing import List, Optional, Union
 
 from alitra import Pose
-from transitions import State
 
-from robot_interface.models.mission.status import RobotStatus
 from robot_interface.models.robots.battery_state import BatteryState
+from robot_interface.models.mission.status import RobotStatus, MissionStatus, TaskStatus
+from robot_interface.models.mission.task import TaskTypes
+from robot_interface.models.exceptions.robot_exceptions import ErrorReason
 
 
 @dataclass
@@ -54,12 +55,7 @@ class DocumentInfo:
 class RobotStatusPayload:
     isar_id: str
     robot_name: str
-    robot_status: RobotStatus
-    previous_robot_status: RobotStatus
-    current_isar_state: State
-    current_mission_id: str
-    current_task_id: str
-    current_step_id: str
+    status: RobotStatus
     timestamp: datetime
 
 
@@ -82,3 +78,38 @@ class RobotHeartbeatPayload:
     isar_id: str
     robot_name: str
     timestamp: datetime
+
+
+@dataclass
+class MissionPayload:
+    isar_id: str
+    robot_name: str
+    mission_id: Optional[str]
+    status: Optional[MissionStatus]
+    error_reason: Optional[ErrorReason]
+    error_description: Optional[str]
+    timestamp: datetime
+
+
+@dataclass
+class TaskPayload:
+    isar_id: str
+    robot_name: str
+    mission_id: Optional[str]
+    task_id: Optional[str]
+    status: Optional[TaskStatus]
+    task_type: Optional[TaskTypes]
+    error_reason: Optional[ErrorReason]
+    error_description: Optional[str]
+    timestamp: datetime
+
+
+@dataclass
+class InspectionResultPayload:
+    isar_id: str
+    robot_name: str
+    inspection_id: str
+    inspection_path: Union[str, dict]
+    installation_code: str
+    analysis_to_be_run: Optional[str]
+    timestamp: datetime