[DPE-2317] Small upgrade fixes (#151)

src/events/password_actions.py

@@ -43,6 +43,12 @@ def _set_password_action(self, event: ActionEvent) -> None:
             event.fail(msg)
             return
 
+        if not self.charm.upgrade.idle:
+            msg = f"Cannot set password while upgrading (upgrade_stack: {self.charm.upgrade.upgrade_stack})"
+            logger.error(msg)
+            event.fail(msg)
+            return
+
         if not self.charm.healthy:
             msg = "Unit is not healthy"
             logger.error(msg)

src/events/upgrade.py

@@ -41,6 +41,15 @@ def __init__(self, charm: "KafkaCharm", **kwargs):
         super().__init__(charm, **kwargs)
         self.charm = charm
 
+    @property
+    def idle(self) -> bool:
+        """Checks if cluster state is idle.
+
+        Returns:
+            True if cluster state is idle. Otherwise False
+        """
+        return not bool(self.upgrade_stack)
+
     @property
     def current_version(self) -> str:
         """Get current Kafka version."""

tests/unit/test_upgrade.py

@@ -79,6 +79,28 @@ def test_build_upgrade_stack(harness: Harness):
     assert len(stack) == len(set(stack))
 
 
+@pytest.mark.parametrize("upgrade_stack", ([], [0]))
+def test_run_password_rotation_while_upgrading(harness, upgrade_stack):
+    harness.charm.upgrade.upgrade_stack = upgrade_stack
+    harness.set_leader(True)
+
+    mock_event = MagicMock()
+    mock_event.params = {"username": "admin"}
+
+    with (
+        patch("charm.KafkaCharm.healthy", new_callable=PropertyMock, return_value=True),
+        patch("managers.auth.AuthManager.add_user"),
+    ):
+        harness.charm.password_action_events._set_password_action(mock_event)
+
+    if not upgrade_stack:
+        mock_event.set_results.assert_called()
+    else:
+        mock_event.fail.assert_called_with(
+            f"Cannot set password while upgrading (upgrade_stack: {upgrade_stack})"
+        )
+
+
 def test_kafka_dependency_model():
     assert sorted(KafkaDependencyModel.__fields__.keys()) == sorted(DEPENDENCIES.keys())