From 48c5e4323e67088115e162cab1be529e08f13053 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Anders=20Chirico=20Indreb=C3=B8?= <ACIN@equinor.com>
Date: Thu, 22 Feb 2024 08:10:55 +0100
Subject: [PATCH] Added Offline and Busy to robot_status

---
 src/isar_exr/api/energy_robotics_api.py | 32 ++++++++++++++++++++++++-
 src/isar_exr/robotinterface.py          |  9 +++++--
 2 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/src/isar_exr/api/energy_robotics_api.py b/src/isar_exr/api/energy_robotics_api.py
index 5954b3f..69ad672 100644
--- a/src/isar_exr/api/energy_robotics_api.py
+++ b/src/isar_exr/api/energy_robotics_api.py
@@ -449,7 +449,7 @@ def wake_up_robot(
         wake_up_robot_mutation.variable_definitions = variable_definitions_graphql
 
         try:
-            result: Dict[str, Any] = self.client.query(
+            result: Dict[str, Any] = self.client.query( # TODO: consider checking if request was accepted
                 dsl_gql(wake_up_robot_mutation), params
             )
         except Exception:
@@ -547,6 +547,36 @@ def get_battery_level(self, exr_robot_id: str) -> Optional[float]:
             "percentage"
         ]
         return battery_level
+    
+    def is_connected(self, exr_robot_id: str) -> bool:
+        params: dict = {"robotID": exr_robot_id}
+
+        variable_definitions_graphql: DSLVariableDefinitions = DSLVariableDefinitions()
+
+        check_is_connected_query: DSLQuery = DSLQuery(
+            self.schema.Query.currentRobotStatus.args(
+                robotID=variable_definitions_graphql.robotID
+            ).select(
+                self.schema.RobotStatusType.isConnected,
+            )
+        )
+
+        check_is_connected_query.variable_definitions = variable_definitions_graphql
+
+        try:
+            result: Dict[str, Any] = self.client.query(
+                dsl_gql(check_is_connected_query), params
+            )
+        except TimeoutError:
+            return False
+        except Exception:
+            message: str = "Could not check robot battery level"
+            self.logger.error(message)
+            raise RobotMissionStatusException(
+                error_description=message,
+            )
+
+        return result["currentRobotStatus"]["isConnected"]
 
     def create_mission_definition(
         self, site_id: str, mission_name: str, robot_id: str
diff --git a/src/isar_exr/robotinterface.py b/src/isar_exr/robotinterface.py
index 85f41a6..44a7698 100644
--- a/src/isar_exr/robotinterface.py
+++ b/src/isar_exr/robotinterface.py
@@ -297,9 +297,14 @@ def get_telemetry_publishers(
         return publisher_threads
 
     def robot_status(self) -> RobotStatus:
-        # TODO: check if robot is running a task, or check if it is awake?
-        # TODO: use currentMissionExecution to see if Busy
         # TODO: find endpoint to check if it is stuck, maybe MissionExecutionStatusEnum.PAUSED
+        if not self.api.is_connected():
+            return RobotStatus.Offline
+        
+        mission_status: MissionStatus = self.api.get_mission_status(settings.ROBOT_EXR_ID)
+        if mission_status == MissionStatus.Paused or mission_status == MissionStatus.NotStarted or \
+            mission_status == MissionStatus.InProgress:
+            return RobotStatus.Busy
         return RobotStatus.Available
 
     def _get_pose_telemetry(self, isar_id: str, robot_name: str) -> str: