From 12830d6ff185e6ec1cb881b9dc4d255200536fe9 Mon Sep 17 00:00:00 2001
From: Mateo Florido <mateo.florido@canonical.com>
Date: Wed, 11 Dec 2024 15:21:41 -0500
Subject: [PATCH] Improve Upgrade Status

---
 charms/worker/k8s/src/charm.py | 24 ++++++++++++++++++------
 1 file changed, 18 insertions(+), 6 deletions(-)

diff --git a/charms/worker/k8s/src/charm.py b/charms/worker/k8s/src/charm.py
index fc438a8b..24935e2c 100755
--- a/charms/worker/k8s/src/charm.py
+++ b/charms/worker/k8s/src/charm.py
@@ -729,10 +729,14 @@ def _announce_kubernetes_version(self):
         if not local_version:
             raise ReconcilerError("k8s-snap is not installed")
 
-        peer = self.model.get_relation(CLUSTER_RELATION)
-        worker = self.model.get_relation(CLUSTER_WORKER_RELATION)
+        relations = (
+            self.model.get_relation(CLUSTER_RELATION),
+            self.model.get_relation(CLUSTER_WORKER_RELATION),
+        )
+        relation_types = ("peer", "worker")
+        waiting_units = {rtype: 0 for rtype in relation_types}
 
-        for relation in (peer, worker):
+        for relation, relation_type in zip(relations, relation_types):
             if not relation:
                 continue
             units = (unit for unit in relation.units if unit.name != self.unit.name)
@@ -741,11 +745,19 @@ def _announce_kubernetes_version(self):
                 if not unit_version:
                     raise ReconcilerError(f"Waiting for version from {unit.name}")
                 if unit_version != local_version:
-                    # NOTE: Add a check to validate if we are doing an upgrade
-                    status.add(ops.WaitingStatus("Upgrading the cluster"))
-                    return
+                    waiting_units[relation_type] += 1
             relation.data[self.app]["version"] = local_version
 
+        if any(waiting_units.values()):
+            waiting_parts = []
+            if waiting_units["peer"]:
+                waiting_parts.append(f"{waiting_units['peer']} Control Plane")
+            if waiting_units["worker"]:
+                waiting_parts.append(f"{waiting_units['worker']} Worker")
+            status_msg = ", ".join(waiting_parts)
+            status.add(ops.WaitingStatus(f"Waiting {status_msg} to upgrade"))
+            raise ReconcilerError("Waiting for all units to upgrade")
+
     def _get_proxy_env(self) -> Dict[str, str]:
         """Retrieve the Juju model config proxy values.