[DPE-4598] Handle upgrade of top of the stack Juju leader (#176)

* Handle upgrade of top of the stack Juju leader Signed-off-by: Marcelo Henrique Neppel <[email protected]> * Add unit tests Signed-off-by: Marcelo Henrique Neppel <[email protected]> --------- Signed-off-by: Marcelo Henrique Neppel <[email protected]>
canonical · Jun 13, 2024 · 887b341 · 887b341
1 parent 5c221f2
commit 887b341
Show file tree

Hide file tree

Showing 2 changed files with 68 additions and 1 deletion.
diff --git a/lib/charms/data_platform_libs/v0/upgrade.py b/lib/charms/data_platform_libs/v0/upgrade.py
@@ -285,7 +285,7 @@ def restart(self, event) -> None:
 
 # Increment this PATCH version before using `charmcraft publish-lib` or reset
 # to 0 if you are raising the major API version
-LIBPATCH = 16
+LIBPATCH = 17
 
 PYDEPS = ["pydantic>=1.10,<2", "poetry-core"]
 
@@ -907,6 +907,17 @@ def _on_upgrade_charm(self, event: UpgradeCharmEvent) -> None:
                     logger.error(e)
                     self.set_unit_failed()
                     return
+                top_unit_id = self.upgrade_stack[-1]
+                top_unit = self.charm.model.get_unit(f"{self.charm.app.name}/{top_unit_id}")
+                if (
+                    top_unit == self.charm.unit
+                    and self.peer_relation.data[self.charm.unit].get("state") == "recovery"
+                ):
+                    # While in a rollback and the Juju leader unit is the top unit in the upgrade stack, emit the event
+                    # for this unit to start the rollback.
+                    self.peer_relation.data[self.charm.unit].update({"state": "ready"})
+                    self.on_upgrade_changed(event)
+                    return
             self.charm.unit.status = WaitingStatus("other units upgrading first...")
             self.peer_relation.data[self.charm.unit].update({"state": "ready"})
 

diff --git a/tests/unit/test_upgrade.py b/tests/unit/test_upgrade.py
@@ -753,6 +753,62 @@ def test_upgrade_charm_runs_checks_on_leader(harness, mocker):
     harness.charm.upgrade._upgrade_supported_check.assert_called_once()
 
 
+@pytest.mark.parametrize(
+    "substrate,leader,call_count,state",
+    [
+        ("vm", True, 1, "ready"),
+        ("vm", False, 0, "ready"),
+        ("k8s", True, 0, "upgrading"),
+        ("k8s", False, 0, "upgrading"),
+    ],
+)
+def test_upgrade_charm_runs_upgrade_changed_on_leader_first_to_rollback(
+    harness, mocker, substrate, leader, call_count, state
+):
+    harness.charm.upgrade = GandalfUpgrade(
+        charm=harness.charm, dependency_model=GandalfModel(**GANDALF_DEPS), substrate=substrate
+    )
+    harness.add_relation("upgrade", "gandalf")
+    harness.set_leader(leader)
+    harness.charm.upgrade.peer_relation.data[harness.charm.unit].update({"state": "recovery"})
+    harness.charm.upgrade.upgrade_stack = [2, 1, 0]
+
+    mocker.patch.object(harness.charm.upgrade, "_upgrade_supported_check")
+    mocker.patch.object(harness.charm.upgrade, "on_upgrade_changed")
+    harness.charm.on.upgrade_charm.emit()
+
+    assert harness.charm.upgrade.on_upgrade_changed.call_count == call_count
+    assert harness.charm.upgrade.state == state
+
+
+@pytest.mark.parametrize(
+    "substrate,initial_state,upgrade_stack,final_state",
+    [
+        ("vm", "idle", [2, 1, 0], "ready"),
+        ("vm", "recovery", [2, 0, 1], "ready"),
+        ("vm", "recovery", [0, 2, 1], "ready"),
+        ("k8s", "recovery", [2, 1, 0], "upgrading"),
+    ],
+)
+def test_upgrade_charm_doesnt_run_upgrade_changed_on_leader_not_first_to_rollback(
+    harness, mocker, substrate, initial_state, upgrade_stack, final_state
+):
+    harness.charm.upgrade = GandalfUpgrade(
+        charm=harness.charm, dependency_model=GandalfModel(**GANDALF_DEPS), substrate=substrate
+    )
+    harness.add_relation("upgrade", "gandalf")
+    harness.set_leader(True)
+    harness.charm.upgrade.peer_relation.data[harness.charm.unit].update({"state": initial_state})
+    harness.charm.upgrade.upgrade_stack = upgrade_stack
+
+    mocker.patch.object(harness.charm.upgrade, "_upgrade_supported_check")
+    mocker.patch.object(harness.charm.upgrade, "on_upgrade_changed")
+    harness.charm.on.upgrade_charm.emit()
+
+    harness.charm.upgrade.on_upgrade_changed.assert_not_called()
+    assert harness.charm.upgrade.state == final_state
+
+
 @pytest.mark.parametrize("substrate,state", [("vm", "ready"), ("k8s", "upgrading")])
 def test_upgrade_charm_sets_right_state(harness, mocker, substrate, state):
     harness.charm.upgrade = GandalfUpgrade(