Fix workload_version in status message for VM upgrade (#128)

## Issue
Leftover from porting k8s implementation: it was assumed that
`workload_version` file matched the currently installed workload. This
is true for k8s, where workload + charm code are upgraded together, but
not true for VM (where charm code upgraded on all units at once, then
workload upgraded by us in charm code)

Impact: incorrect workload version in status message if snap not
upgraded

## Solution
Save `workload_version` value in unit databag when saving snap revision

src/machine_upgrade.py

@@ -26,8 +26,8 @@ class Upgrade(upgrade.Upgrade):
     @property
     def unit_state(self) -> typing.Optional[str]:
         if (
-            self._unit_workload_version is not None
-            and self._unit_workload_version != self._app_workload_version
+            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"
@@ -37,28 +37,31 @@ def unit_state(self) -> typing.Optional[str]:
     def unit_state(self, value: str) -> None:
         if value == "healthy":
             # Set snap revision on first install
-            self._unit_databag["snap_revision"] = snap.REVISION
-            logger.debug(f"Saved {snap.REVISION=} in unit databag while setting state healthy")
+            self._unit_workload_container_version = snap.REVISION
+            self._unit_workload_version = self._current_versions["workload"]
+            logger.debug(
+                f'Saved {snap.REVISION=} and {self._current_versions["workload"]=} in unit databag while setting state healthy'
+            )
         # Super call
         upgrade.Upgrade.unit_state.fset(self, value)
 
     def _get_unit_healthy_status(
         self, *, workload_status: typing.Optional[ops.StatusBase]
     ) -> ops.StatusBase:
-        if self._unit_workload_version == self._app_workload_version:
+        if self._unit_workload_container_version == self._app_workload_container_version:
             if isinstance(workload_status, ops.WaitingStatus):
                 return ops.WaitingStatus(
-                    f'Router {self._current_versions["workload"]} rev {self._unit_workload_version}'
+                    f"Router {self._unit_workload_version} rev {self._unit_workload_container_version}"
                 )
             return ops.ActiveStatus(
-                f'Router {self._current_versions["workload"]} rev {self._unit_workload_version} running'
+                f"Router {self._unit_workload_version} rev {self._unit_workload_container_version} running"
             )
         if isinstance(workload_status, ops.WaitingStatus):
             return ops.WaitingStatus(
-                f'Charmed operator upgraded. Router {self._current_versions["workload"]} rev {self._unit_workload_version}'
+                f"Charmed operator upgraded. Router {self._unit_workload_version} rev {self._unit_workload_container_version}"
             )
         return ops.WaitingStatus(
-            f'Charmed operator upgraded. Router {self._current_versions["workload"]} rev {self._unit_workload_version} running'
+            f"Charmed operator upgraded. Router {self._unit_workload_version} rev {self._unit_workload_container_version} running"
         )
 
     @property
@@ -73,7 +76,7 @@ def app_status(self) -> typing.Optional[ops.StatusBase]:
         return super().app_status
 
     @property
-    def _unit_workload_versions(self) -> typing.Dict[str, str]:
+    def _unit_workload_container_versions(self) -> typing.Dict[str, str]:
         """{Unit name: installed snap revision}"""
         versions = {}
         for unit in self._sorted_units:
@@ -82,15 +85,28 @@ def _unit_workload_versions(self) -> typing.Dict[str, str]:
         return versions
 
     @property
-    def _unit_workload_version(self) -> typing.Optional[str]:
+    def _unit_workload_container_version(self) -> typing.Optional[str]:
         """Installed snap revision for this unit"""
         return self._unit_databag.get("snap_revision")
 
+    @_unit_workload_container_version.setter
+    def _unit_workload_container_version(self, value: str):
+        self._unit_databag["snap_revision"] = value
+
     @property
-    def _app_workload_version(self) -> str:
+    def _app_workload_container_version(self) -> str:
         """Snap revision for current charm code"""
         return snap.REVISION
 
+    @property
+    def _unit_workload_version(self) -> typing.Optional[str]:
+        """Installed MySQL Router version for this unit"""
+        return self._unit_databag.get("workload_version")
+
+    @_unit_workload_version.setter
+    def _unit_workload_version(self, value: str):
+        self._unit_databag["workload_version"] = value
+
     def reconcile_partition(self, *, action_event: ops.ActionEvent = None) -> None:
         """Handle Juju action to confirm first upgraded unit is healthy and resume upgrade."""
         if action_event:
@@ -118,7 +134,7 @@ def upgrade_resumed(self, value: bool):
 
     @property
     def authorized(self) -> bool:
-        assert self._unit_workload_version != self._app_workload_version
+        assert self._unit_workload_container_version != self._app_workload_container_version
         for index, unit in enumerate(self._sorted_units):
             if unit.name == self._unit.name:
                 # Higher number units have already upgraded
@@ -128,7 +144,8 @@ def authorized(self) -> bool:
                     return self.upgrade_resumed
                 return True
             if (
-                self._unit_workload_versions.get(unit.name) != self._app_workload_version
+                self._unit_workload_container_versions.get(unit.name)
+                != self._app_workload_container_version
                 or self._peer_relation.data[unit].get("state") != "healthy"
             ):
                 # Waiting for higher number units to upgrade
@@ -139,5 +156,8 @@ def upgrade_unit(self, *, workload_: workload.Workload, tls: bool) -> None:
         logger.debug(f"Upgrading {self.authorized=}")
         self.unit_state = "upgrading"
         workload_.upgrade(unit=self._unit, tls=tls)
-        self._unit_databag["snap_revision"] = snap.REVISION
-        logger.debug(f"Saved {snap.REVISION=} in unit databag after upgrade")
+        self._unit_workload_container_version = snap.REVISION
+        self._unit_workload_version = self._current_versions["workload"]
+        logger.debug(
+            f'Saved {snap.REVISION=} and {self._current_versions["workload"]=} in unit databag after upgrade'
+        )

src/upgrade.py

@@ -113,10 +113,12 @@ def is_compatible(self) -> bool:
 
     @property
     def in_progress(self) -> bool:
-        logger.debug(f"{self._app_workload_version=} {self._unit_workload_versions=}")
+        logger.debug(
+            f"{self._app_workload_container_version=} {self._unit_workload_container_versions=}"
+        )
         return any(
-            version != self._app_workload_version
-            for version in self._unit_workload_versions.values()
+            version != self._app_workload_container_version
+            for version in self._unit_workload_container_versions.values()
         )
 
     @property
@@ -179,28 +181,28 @@ def upgrade_resumed(self) -> bool:
 
     @property
     @abc.abstractmethod
-    def _unit_workload_versions(self) -> typing.Dict[str, str]:
-        """{Unit name: unique identifier for unit's workload version}
+    def _unit_workload_container_versions(self) -> typing.Dict[str, str]:
+        """{Unit name: unique identifier for unit's workload container version}
 
         If and only if this version changes, the workload will restart (during upgrade or
         rollback).
 
         On Kubernetes, the workload & charm are upgraded together
         On machines, the charm is upgraded before the workload
 
-        This identifier should be comparable to `_app_workload_version` to determine if the unit &
-        app are the same workload version.
+        This identifier should be comparable to `_app_workload_container_version` to determine if
+        the unit & app are the same workload container version.
         """
 
     @property
     @abc.abstractmethod
-    def _app_workload_version(self) -> str:
-        """Unique identifier for the app's workload version
+    def _app_workload_container_version(self) -> str:
+        """Unique identifier for the app's workload container version
 
         This should match the workload version in the current Juju app charm version.
 
-        This identifier should be comparable to `_get_unit_workload_version` to determine if the
-        app & unit are the same workload version.
+        This identifier should be comparable to `_unit_workload_container_versions` to determine if
+        the app & unit are the same workload container version.
         """
 
     @abc.abstractmethod