From 6602c5cdca858256914930f2515a01bf12882a7f Mon Sep 17 00:00:00 2001
From: "Mariana R. Santos" <mrica@equinor.com>
Date: Tue, 24 Oct 2023 17:01:46 +0200
Subject: [PATCH] Add function to check if pipeline is completed

---
 src/isar_exr/api/energy_robotics_api.py | 39 +++++++++++++++++++++++++
 src/isar_exr/config/settings.py         |  3 ++
 src/isar_exr/robotinterface.py          |  6 +++-
 3 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/src/isar_exr/api/energy_robotics_api.py b/src/isar_exr/api/energy_robotics_api.py
index 4435693..772a7d8 100644
--- a/src/isar_exr/api/energy_robotics_api.py
+++ b/src/isar_exr/api/energy_robotics_api.py
@@ -61,6 +61,10 @@ def get_mission_status(self, exr_robot_id: str) -> MissionStatus:
         response_dict: dict[str, Any] = self.client.query(
             dsl_gql(current_mission_execution_query), params
         )
+        
+        if response_dict["currentMissionExecution"] is None:
+            raise NoMissionRunningException
+        
         step_status = ExrMissionStatus(
             response_dict["currentMissionExecution"]["status"]
         )
@@ -660,6 +664,36 @@ def commit_site_to_snapshot(self, stage_id: str) -> str:
             )
 
         return response_dict["commitSiteChanges"]["id"]
+    
+    def is_pipeline_completed(self, site_id: str) -> bool:
+
+        variable_definitions_graphql: DSLVariableDefinitions = DSLVariableDefinitions()
+
+        current_processing_pipeline: DSLQuery = DSLQuery(
+            self.schema.Query.currentSiteSnapshotHeadSelectionProcessingPipeline.args(
+                siteId=variable_definitions_graphql.siteId
+            ).select(self.schema.ProcessingPipelineType.stages.select(self.schema.ProcessingPipelineStageType.state))
+        )
+
+        current_processing_pipeline.variable_definitions = variable_definitions_graphql
+
+        params: dict = {"siteId": site_id}
+
+        try:
+            response_dict: dict[str, Any] = self.client.query(
+                dsl_gql(current_processing_pipeline), params
+            )
+        except Exception as e:
+            message: str = "Could not get current processing pipeline"
+            self.logger.error(message)
+            raise RobotAPIException(
+                error_description=message,
+            )
+        
+        if response_dict["currentSiteSnapshotHeadSelectionProcessingPipeline"]["stages"][0]["state"] == "COMPLETED":
+            return True
+        return False
+
 
     def set_snapshot_as_head(self, snapshot_id: str, site_id: str) -> str:
         params: dict[str, Any] = {"siteId": site_id, "siteSnapshotId": snapshot_id}
@@ -716,3 +750,8 @@ def get_current_site_stage(self, site_id: str) -> str:
         if response_dict["currentSiteStage"] is not None:
             return response_dict["currentSiteStage"]["id"]
         return None
+
+if __name__=="__main__":
+    api: EnergyRoboticsApi = EnergyRoboticsApi()
+    exr_robot_id: str = settings.ROBOT_EXR_ID
+    api.get_status_current_processing_pipeline(settings.ROBOT_EXR_SITE_ID)
diff --git a/src/isar_exr/config/settings.py b/src/isar_exr/config/settings.py
index b88e57e..fdea0d3 100644
--- a/src/isar_exr/config/settings.py
+++ b/src/isar_exr/config/settings.py
@@ -44,6 +44,9 @@ def __init__(self) -> None:
         "../../../docs/schema.graphql"
     )
 
+    # API sleep time
+    API_SLEEP_TIME: int = Field(default=1)
+
     model_config = SettingsConfigDict(
         env_prefix="EXR_",
         env_file_encoding="utf-8",
diff --git a/src/isar_exr/robotinterface.py b/src/isar_exr/robotinterface.py
index 354a8c5..981e38b 100644
--- a/src/isar_exr/robotinterface.py
+++ b/src/isar_exr/robotinterface.py
@@ -103,10 +103,14 @@ def initiate_mission(self, mission: Mission) -> None:
                     poi_ids.append(poi_id)
 
         snapshot_id: str = self.api.commit_site_to_snapshot(stage_id=stage_id)
+
         self.api.set_snapshot_as_head(
             snapshot_id=snapshot_id, site_id=settings.ROBOT_EXR_SITE_ID
         )
 
+        while (not self.api.is_pipeline_completed(site_id=settings.ROBOT_EXR_SITE_ID)):
+            time.sleep(settings.API_SLEEP_TIME)
+
         mission_definition_id: str = self.api.create_mission_definition(
             site_id=settings.ROBOT_EXR_SITE_ID,
             mission_name=mission.id,
@@ -372,4 +376,4 @@ def _add_dock_robot_task_to_mission(
         self.api.add_task_to_mission_definition(
             task_id=dock_task_id,
             mission_definition_id=mission_definition_id,
-        )
+        )
\ No newline at end of file