Improve Upgrade Status

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.