From e0ef5edab026c636c30f3b88cfc14f4b17b31d82 Mon Sep 17 00:00:00 2001
From: Smruti Ranjan Senapati <smrtsenapati59@gmail.com>
Date: Wed, 17 Jul 2024 10:04:54 +0530
Subject: [PATCH] fix(deployments): fix deployment wait command

---
 riocli/deployment/model.py |  1 +
 riocli/deployment/wait.py  |  8 ++++++--
 riocli/network/delete.py   |  2 ++
 riocli/v2client/client.py  | 12 +++++-------
 riocli/v2client/error.py   |  4 ++++
 5 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/riocli/deployment/model.py b/riocli/deployment/model.py
index f3a4ff45..693fa51c 100644
--- a/riocli/deployment/model.py
+++ b/riocli/deployment/model.py
@@ -16,6 +16,7 @@
 from munch import unmunchify
 from rapyuta_io import Client
 from rapyuta_io.clients.catalog_client import Package
+from rapyuta_io.clients.deployment import DeploymentNotRunningException
 from rapyuta_io.clients.package import ProvisionConfiguration, RestartPolicy
 from rapyuta_io.clients.rosbag import (OverrideOptions, ROSBagCompression, ROSBagJob, ROSBagOnDemandUploadOptions,
                                        ROSBagOptions, ROSBagTimeRange, ROSBagUploadTypes, TopicOverrideInfo,
diff --git a/riocli/deployment/wait.py b/riocli/deployment/wait.py
index 1a0c7d8c..5579ee2c 100644
--- a/riocli/deployment/wait.py
+++ b/riocli/deployment/wait.py
@@ -17,7 +17,7 @@
 from riocli.config import new_v2_client
 from riocli.constants import Colors, Symbols
 from riocli.utils.spinner import with_spinner
-from riocli.v2client.error import RetriesExhausted, DeploymentNotRunning
+from riocli.v2client.error import RetriesExhausted, DeploymentNotRunning, ImagePullError
 
 
 @click.command(
@@ -38,7 +38,7 @@ def wait_for_deployment(
     try:
         client = new_v2_client()
         deployment = client.poll_deployment(deployment_name)
-        spinner.text = click.style('Deployment status: {}'.format(deployment.status.status), fg=Colors.GREEN)
+        spinner.text = click.style('Phase: Succeeded Status: {}'.format(deployment.status.status), fg=Colors.GREEN)
         spinner.green.ok(Symbols.SUCCESS)
     except RetriesExhausted as e:
         spinner.write(click.style(str(e), fg=Colors.RED))
@@ -48,6 +48,10 @@ def wait_for_deployment(
         spinner.text = click.style(str(e), fg=Colors.RED)
         spinner.red.fail(Symbols.ERROR)
         raise SystemExit(1)
+    except ImagePullError as e:
+        spinner.text = click.style(str(e), fg=Colors.RED)
+        spinner.red.fail(Symbols.ERROR)
+        raise SystemExit(1)
     except Exception as e:
         spinner.text = click.style(str(e), fg=Colors.RED)
         spinner.red.fail(Symbols.ERROR)
diff --git a/riocli/network/delete.py b/riocli/network/delete.py
index ed721fb2..2ea7fc68 100644
--- a/riocli/network/delete.py
+++ b/riocli/network/delete.py
@@ -27,6 +27,7 @@
 from rapyuta_io import Client
 from riocli.network.model import Network
 
+
 @click.command(
     'delete',
     cls=HelpColorsCommand,
@@ -112,6 +113,7 @@ def delete_network(
         spinner.red.fail(Symbols.ERROR)
         raise SystemExit(1) from e
 
+
 def _apply_delete(client: Client, result: Queue, network: Network) -> None:
     try:
         client.delete_network(network_name=network.metadata.name)
diff --git a/riocli/v2client/client.py b/riocli/v2client/client.py
index 85b6bb04..a035aa66 100644
--- a/riocli/v2client/client.py
+++ b/riocli/v2client/client.py
@@ -27,7 +27,7 @@
 from rapyuta_io.utils.rest_client import HttpMethod, RestClient
 
 from riocli.v2client.enums import DeploymentPhaseConstants
-from riocli.v2client.error import RetriesExhausted, DeploymentNotRunning
+from riocli.v2client.error import RetriesExhausted, DeploymentNotRunning, ImagePullError
 
 
 class DeploymentNotFound(Exception):
@@ -1048,15 +1048,13 @@ def poll_deployment(
                 return deployment
 
             if status.phase == DeploymentPhaseConstants.DeploymentPhaseProvisioning.value:
-                errors = status.error_codes or []
+                errors = status.get('error_codes', [])
                 if 'DEP_E153' in errors:  # DEP_E153 (image-pull error) will persist across retries
-                    return deployment
+                    raise ImagePullError('Deployment not running. Phase: Provisioning Status: {}'.format(status.phase))
             elif status.phase == DeploymentPhaseConstants.DeploymentPhaseSucceeded.value:
                 return deployment
-            elif status.phase in [DeploymentPhaseConstants.DeploymentPhaseFailedToUpdate.value,
-                                  DeploymentPhaseConstants.DeploymentPhaseFailedToStart.value,
-                                  DeploymentPhaseConstants.DeploymentPhaseStopped.value]:
-                raise DeploymentNotRunning('Deployment not running. Deployment status: {}'.format(status.phase))
+            elif status.phase == DeploymentPhaseConstants.DeploymentPhaseStopped.value:
+                raise DeploymentNotRunning('Deployment not running. Phase: Stopped  Status: {}'.format(status.phase))
 
             time.sleep(sleep_interval)
             deployment = self.get_deployment(name)
diff --git a/riocli/v2client/error.py b/riocli/v2client/error.py
index 98881e67..bef94764 100644
--- a/riocli/v2client/error.py
+++ b/riocli/v2client/error.py
@@ -5,5 +5,9 @@ def __init__(self, msg=None):
 
 
 class DeploymentNotRunning(Exception):
+    def __init__(self, msg=None):
+        Exception.__init__(self, msg)
+
+class ImagePullError(Exception):
     def __init__(self, msg=None):
         Exception.__init__(self, msg)
\ No newline at end of file