Add upgrade UnitState enum (#134)

Author: carlcsaposs-canonical

src/abstract_charm.py

@@ -254,7 +254,7 @@ def reconcile(self, event=None) -> None:  # noqa: C901
             # Run before checking `self._upgrade.is_compatible` in case incompatible upgrade was
             # forced & completed on all units.
             self._upgrade.set_versions_in_app_databag()
-        if self._upgrade.unit_state == "restarting":  # Kubernetes only
+        if self._upgrade.unit_state is upgrade.UnitState.RESTARTING:  # Kubernetes only
             if not self._upgrade.is_compatible:
                 logger.info(
                     "Upgrade incompatible. If you accept potential *data loss* and *downtime*, you can continue with `resume-upgrade force=true`"
@@ -268,7 +268,7 @@ def reconcile(self, event=None) -> None:  # noqa: C901
             if not self._upgrade.is_compatible:
                 self.set_status(event=event)
                 return
-            if self._upgrade.unit_state == "outdated":
+            if self._upgrade.unit_state is upgrade.UnitState.OUTDATED:
                 if self._upgrade.authorized:
                     self._upgrade.upgrade_unit(
                         event=event,
@@ -330,7 +330,7 @@ def reconcile(self, event=None) -> None:  # noqa: C901
             # Empty waiting status means we're waiting for database requires relation before
             # starting workload
             if not workload_.status or workload_.status == ops.WaitingStatus():
-                self._upgrade.unit_state = "healthy"
+                self._upgrade.unit_state = upgrade.UnitState.HEALTHY
             if self._unit_lifecycle.authorized_leader:
                 self._upgrade.reconcile_partition()
             self.set_status(event=event)

src/machine_charm.py

@@ -165,7 +165,7 @@ def _on_force_upgrade_action(self, event: ops.ActionEvent) -> None:
             logger.debug(f"Force upgrade event failed: {message}")
             event.fail(message)
             return
-        if self._upgrade.unit_state != "outdated":
+        if self._upgrade.unit_state is not upgrade.UnitState.OUTDATED:
             message = "Unit already upgraded"
             logger.debug(f"Force upgrade event failed: {message}")
             event.fail(message)

src/machine_upgrade.py

@@ -27,18 +27,18 @@ class Upgrade(upgrade.Upgrade):
     """In-place upgrades on machines"""
 
     @property
-    def unit_state(self) -> typing.Optional[str]:
+    def unit_state(self) -> typing.Optional[upgrade.UnitState]:
         if (
             self._unit_workload_container_version is not None
             and self._unit_workload_container_version != self._app_workload_container_version
         ):
             logger.debug("Unit upgrade state: outdated")
-            return "outdated"
+            return upgrade.UnitState.OUTDATED
         return super().unit_state
 
     @unit_state.setter
-    def unit_state(self, value: str) -> None:
-        if value == "healthy":
+    def unit_state(self, value: upgrade.UnitState) -> None:
+        if value is upgrade.UnitState.HEALTHY:
             # Set snap revision on first install
             self._unit_workload_container_version = snap.REVISION
             self._unit_workload_version = self._current_versions["workload"]
@@ -146,10 +146,13 @@ def authorized(self) -> bool:
                     logger.debug(f"Second unit authorized to upgrade if {self.upgrade_resumed=}")
                     return self.upgrade_resumed
                 return True
+            state = self._peer_relation.data[unit].get("state")
+            if state:
+                state = upgrade.UnitState(state)
             if (
                 self._unit_workload_container_versions.get(unit.name)
                 != self._app_workload_container_version
-                or self._peer_relation.data[unit].get("state") != "healthy"
+                or state is not upgrade.UnitState.HEALTHY
             ):
                 # Waiting for higher number units to upgrade
                 return False
@@ -164,7 +167,7 @@ def upgrade_unit(
         exporter_config: "relations.cos.ExporterConfig",
     ) -> None:
         logger.debug(f"Upgrading {self.authorized=}")
-        self.unit_state = "upgrading"
+        self.unit_state = upgrade.UnitState.UPGRADING
         workload_.upgrade(event=event, unit=self._unit, tls=tls, exporter_config=exporter_config)
         self._unit_workload_container_version = snap.REVISION
         self._unit_workload_version = self._current_versions["workload"]

src/upgrade.py

@@ -9,6 +9,7 @@
 
 import abc
 import copy
+import enum
 import json
 import logging
 import pathlib
@@ -34,6 +35,15 @@ class PeerRelationNotReady(Exception):
     """Upgrade peer relation not available (to this unit)"""
 
 
+class UnitState(str, enum.Enum):
+    """Unit upgrade state"""
+
+    HEALTHY = "healthy"
+    RESTARTING = "restarting"  # Kubernetes only
+    UPGRADING = "upgrading"  # Machines only
+    OUTDATED = "outdated"  # Machines only
+
+
 class Upgrade(abc.ABC):
     """In-place upgrades"""
 
@@ -55,13 +65,14 @@ def __init__(self, charm_: ops.CharmBase) -> None:
             self._current_versions[version] = pathlib.Path(file_name).read_text().strip()
 
     @property
-    def unit_state(self) -> typing.Optional[str]:
+    def unit_state(self) -> typing.Optional[UnitState]:
         """Unit upgrade state"""
-        return self._unit_databag.get("state")
+        if state := self._unit_databag.get("state"):
+            return UnitState(state)
 
     @unit_state.setter
-    def unit_state(self, value: str) -> None:
-        self._unit_databag["state"] = value
+    def unit_state(self, value: UnitState) -> None:
+        self._unit_databag["state"] = value.value
 
     @property
     def is_compatible(self) -> bool: