From 98968a4d6b7c180fd413a6f8f42888383e5a8464 Mon Sep 17 00:00:00 2001
From: Thomas Carmet <8408330+tcarmet@users.noreply.github.com>
Date: Mon, 22 Apr 2024 17:04:25 +0200
Subject: [PATCH] Improvements to list method on aws and gcp (#583)

---
 runner_manager/backend/aws.py    | 19 +++++++++++--------
 runner_manager/backend/gcloud.py | 22 +++++++++++++++-------
 2 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/runner_manager/backend/aws.py b/runner_manager/backend/aws.py
index 420b39e6..8f3f53f8 100644
--- a/runner_manager/backend/aws.py
+++ b/runner_manager/backend/aws.py
@@ -55,13 +55,14 @@ def delete(self, runner: Runner):
                     raise e
         return super().delete(runner)
 
-    def _get_instance_name(self, instance: InstanceTypeDef) -> str:
-        """Get the instance name."""
-        name = ""
+    def _get_tag_value(
+        self, instance: InstanceTypeDef, key: str, default: str = ""
+    ) -> str:
+        """Get the tag value."""
         for tag in instance.get("Tags", []):
-            if tag.get("Key") == "Name":
-                name = tag.get("Value", "")
-        return name
+            if tag.get("Key") == key:
+                return tag.get("Value", default)
+        return default
 
     def list(self) -> List[Runner]:
         """List runners."""
@@ -81,7 +82,7 @@ def list(self) -> List[Runner]:
         for reservation in reservations:
             for instance in reservation.get("Instances", []):
                 instance_id = instance.get("InstanceId", "")
-                name = self._get_instance_name(instance)
+                name = self._get_tag_value(instance, "Name", instance_id)
                 try:
                     runner = Runner.find(
                         Runner.instance_id == instance_id,
@@ -91,8 +92,10 @@ def list(self) -> List[Runner]:
                         name=name,
                         instance_id=instance_id,
                         runner_group_name=self.runner_group,
-                        busy=False,
+                        busy=bool(self._get_tag_value(instance, "busy")),
+                        status=self._get_tag_value(instance, "status", "online"),
                         created_at=instance.get("LaunchTime"),
+                        started_at=instance.get("LaunchTime"),
                     )
                 runners.append(runner)
         return runners
diff --git a/runner_manager/backend/gcloud.py b/runner_manager/backend/gcloud.py
index 6bbf81c3..72dffe46 100644
--- a/runner_manager/backend/gcloud.py
+++ b/runner_manager/backend/gcloud.py
@@ -25,6 +25,7 @@
     ZoneOperationsClient,
 )
 from pydantic import Field
+from redis_om import NotFoundError
 
 from runner_manager.backend.base import BaseBackend
 from runner_manager.models.backend import Backends, GCPConfig, GCPInstanceConfig
@@ -256,13 +257,20 @@ def list(self) -> List[Runner]:
                 manager = instance.labels.get("manager", "")
                 runner_group = instance.labels.get("runner_group", "")
                 if manager == self.manager and runner_group == self.runner_group:
-                    runner: Runner = Runner(
-                        name=instance.name,
-                        instance_id=instance.name,
-                        runner_group_name=self.runner_group,
-                        busy=bool(instance.labels.get("busy", False)),
-                        created_at=instance.creation_timestamp,
-                    )
+                    try:
+                        runner = Runner.find(
+                            Runner.name == instance.name,
+                        ).first()
+                    except NotFoundError:
+                        runner: Runner = Runner(
+                            name=instance.name,
+                            instance_id=instance.name,
+                            runner_group_name=self.runner_group,
+                            busy=bool(instance.labels.get("busy", False)),
+                            status=instance.labels.get("status", "online"),
+                            created_at=instance.creation_timestamp,
+                            started_at=instance.creation_timestamp,
+                        )
                     runners.append(runner)
 
         except Exception as e: