feat: support enable-auto-restarts=False for rabbitmq-server. (#413)

Add support for [`enable-auto-restarts=False`](https://charmhub.io/rabbitmq-server/configuration#enable-auto-restarts)
for rabbitmq-server.

cou/apps/auxiliary.py

@@ -21,7 +21,7 @@
 from cou.apps.base import LONG_IDLE_TIMEOUT, OpenStackApplication
 from cou.apps.factory import AppFactory
 from cou.exceptions import ApplicationError
-from cou.steps import ApplicationUpgradePlan, PreUpgradeStep
+from cou.steps import ApplicationUpgradePlan, PostUpgradeStep, PreUpgradeStep
 from cou.utils.app_utils import set_require_osd_release_option
 from cou.utils.juju_utils import Unit
 from cou.utils.openstack import (
@@ -183,6 +183,103 @@ class RabbitMQServer(AuxiliaryApplication):
     # COU changes to 3.9 if the channel is set to 3.8
     multiple_channels = True
 
+    def pre_upgrade_steps(
+        self, target: OpenStackRelease, units: Optional[list[Unit]]
+    ) -> list[PreUpgradeStep]:
+        """Pre Upgrade steps planning.
+
+        :param target: OpenStack release as target to upgrade.
+        :type target: OpenStackRelease
+        :param units: Units to generate upgrade plan
+        :type units: Optional[list[Unit]]
+        :return: List of pre upgrade steps.
+        :rtype: list[PreUpgradeStep]
+        """
+        steps = super().pre_upgrade_steps(target, units)
+        # Since auto restart is disabled, we don't know the if the service
+        # has pending events or not, so we want to run `run-deferred-hooks`
+        # action to clear the events before performing upgrade. If there
+        # are no pending events, this step should be a no-op, so it's safe
+        # to run anyway
+        if not self.config["enable-auto-restarts"].get("value"):
+            # Run any deferred events and restart the service. See
+            # https://charmhub.io/rabbitmq-server/actions#run-deferred-hooks
+            units_to_run_action = self.units.values() if units is None else units
+            steps += [
+                PreUpgradeStep(
+                    description="Auto restarts is disabled, will"
+                    f" execute run-deferred-hooks for unit: '{unit.name}'",
+                    coro=self.model.run_action(
+                        unit.name, "run-deferred-hooks", raise_on_failure=True
+                    ),
+                )
+                for unit in units_to_run_action
+            ]
+            steps += [
+                PreUpgradeStep(
+                    description=(
+                        f"Wait for up to {self.wait_timeout}s for app '{self.name}'"
+                        " to reach the idle state"
+                    ),
+                    parallel=False,
+                    coro=self.model.wait_for_active_idle(self.wait_timeout, apps=[self.name]),
+                )
+            ]
+        return steps
+
+    def post_upgrade_steps(
+        self, target: OpenStackRelease, units: Optional[list[Unit]]
+    ) -> list[PostUpgradeStep]:
+        """Post Upgrade steps planning.
+
+        Wait until the application reaches the idle state and then check the target workload.
+
+        :param target: OpenStack release as target to upgrade.
+        :type target: OpenStackRelease
+        :param units: Units to generate post upgrade plan
+        :type units: Optional[list[Unit]]
+        :return: List of post upgrade steps.
+        :rtype: list[PostUpgradeStep]
+        """
+        steps = []
+        # Since the auto restart is disabled, we need to run the
+        # `run-deferred-hooks` action, and restart the service after the
+        # upgrade.
+        if not self.config["enable-auto-restarts"].get("value"):
+            steps += [
+                PostUpgradeStep(
+                    description=(
+                        f"Wait for up to {self.wait_timeout}s for app '{self.name}'"
+                        " to reach the idle state"
+                    ),
+                    parallel=False,
+                    coro=self.model.wait_for_active_idle(self.wait_timeout, apps=[self.name]),
+                )
+            ]
+            # Run any deferred events and restart the service. See
+            # https://charmhub.io/rabbitmq-server/actions#run-deferred-hooks
+            units_to_run_action = self.units.values() if units is None else units
+            steps += [
+                PostUpgradeStep(
+                    description="Auto restarts is disabled, will"
+                    f" execute run-deferred-hooks for unit: '{unit.name}'",
+                    coro=self.model.run_action(
+                        unit.name, "run-deferred-hooks", raise_on_failure=True
+                    ),
+                )
+                for unit in units_to_run_action
+            ]
+        steps += super().post_upgrade_steps(target, units)
+        return steps
+
+    def _check_auto_restarts(self) -> None:
+        """No-op, skip check auto restarts option.
+
+        This method override the parent class's `_check_auto_restarts()` method
+        because the parent class's will raise an `ApplicationError` if
+        `enable-auto-restarts` is `True`.
+        """
+
 
 @AppFactory.register_application(["ceph-mon"])
 class CephMon(AuxiliaryApplication):

tests/unit/apps/test_auxiliary.py

@@ -123,7 +123,10 @@ def test_auxiliary_upgrade_plan_ussuri_to_victoria_change_channel(model):
         can_upgrade_to="3.9/stable",
         charm="rabbitmq-server",
         channel="3.8/stable",
-        config={"source": {"value": "distro"}},
+        config={
+            "source": {"value": "distro"},
+            "enable-auto-restarts": {"value": True},
+        },
         machines=machines,
         model=model,
         origin="ch",
@@ -201,7 +204,10 @@ def test_auxiliary_upgrade_plan_ussuri_to_victoria(model):
         can_upgrade_to="3.9/stable",
         charm="rabbitmq-server",
         channel="3.9/stable",
-        config={"source": {"value": "distro"}},
+        config={
+            "source": {"value": "distro"},
+            "enable-auto-restarts": {"value": True},
+        },
         machines=machines,
         model=model,
         origin="ch",
@@ -272,7 +278,10 @@ def test_auxiliary_upgrade_plan_ussuri_to_victoria_ch_migration(model):
         can_upgrade_to="cs:rabbitmq-server",
         charm="rabbitmq-server",
         channel="stable",
-        config={"source": {"value": "distro"}},
+        config={
+            "source": {"value": "distro"},
+            "enable-auto-restarts": {"value": True},
+        },
         machines=machines,
         model=model,
         origin="cs",
@@ -333,6 +342,110 @@ def test_auxiliary_upgrade_plan_ussuri_to_victoria_ch_migration(model):
     assert_steps(upgrade_plan, expected_plan)
 
 
+def test_rabbitmq_server_upgrade_plan_ussuri_to_victoria_auto_restart_False(model):
+    """Test rabbitmq server upgrade plan from Ussuri to Victoria with enable_auto_restart=False."""
+    target = OpenStackRelease("victoria")
+    machines = {"0": generate_cou_machine("0", "az-0")}
+    app = RabbitMQServer(
+        name="rabbitmq-server",
+        can_upgrade_to="3.9/stable",
+        charm="rabbitmq-server",
+        channel="3.9/stable",
+        config={
+            "source": {"value": "distro"},
+            "enable-auto-restarts": {"value": False},
+        },
+        machines=machines,
+        model=model,
+        origin="ch",
+        series="focal",
+        subordinate_to=[],
+        units={
+            "rabbitmq-server/0": Unit(
+                name="rabbitmq-server/0",
+                workload_version="3.9",
+                machine=machines["0"],
+            ),
+        },
+        workload_version="3.9",
+    )
+
+    expected_plan = ApplicationUpgradePlan(f"Upgrade plan for '{app.name}' to '{target}'")
+    upgrade_packages = PreUpgradeStep(
+        description=f"Upgrade software packages of '{app.name}' from the current APT repositories",
+        parallel=True,
+    )
+    upgrade_packages.add_steps(
+        UnitUpgradeStep(
+            description=f"Upgrade software packages on unit '{unit.name}'",
+            coro=app_utils.upgrade_packages(unit.name, model, None),
+        )
+        for unit in app.units.values()
+    )
+
+    upgrade_steps = [
+        upgrade_packages,
+        PreUpgradeStep(
+            description=f"Refresh '{app.name}' to the latest revision of '3.9/stable'",
+            coro=model.upgrade_charm(app.name, "3.9/stable"),
+        ),
+    ]
+    upgrade_steps += [
+        PreUpgradeStep(
+            description="Auto restarts is disabled, will"
+            f" execute run-deferred-hooks for unit: '{unit.name}'",
+            coro=model.run_action(unit.name, "run-deferred-hooks", raise_on_failure=True),
+        )
+        for unit in app.units.values()
+    ]
+    upgrade_steps += [
+        PreUpgradeStep(
+            description=f"Wait for up to 2400s for app '{app.name}' to reach the idle state",
+            parallel=False,
+            coro=model.wait_for_active_idle(2400, apps=[app.name]),
+        ),
+        UpgradeStep(
+            description=f"Change charm config of '{app.name}' "
+            f"'{app.origin_setting}' to 'cloud:focal-victoria'",
+            parallel=False,
+            coro=model.set_application_config(
+                app.name, {f"{app.origin_setting}": "cloud:focal-victoria"}
+            ),
+        ),
+        PostUpgradeStep(
+            description=f"Wait for up to 2400s for app '{app.name}' to reach the idle state",
+            parallel=False,
+            coro=model.wait_for_active_idle(2400, apps=[app.name]),
+        ),
+    ]
+    upgrade_steps += [
+        PostUpgradeStep(
+            description="Auto restarts is disabled, will"
+            f" execute run-deferred-hooks for unit: '{unit.name}'",
+            coro=model.run_action(unit.name, "run-deferred-hooks", raise_on_failure=True),
+        )
+        for unit in app.units.values()
+    ]
+    upgrade_steps += [
+        PostUpgradeStep(
+            description=f"Wait for up to 2400s for model '{model.name}' to reach the idle state",
+            parallel=False,
+            coro=model.wait_for_active_idle(2400, apps=None),
+        ),
+        PostUpgradeStep(
+            description=f"Verify that the workload of '{app.name}' has been upgraded on units: "
+            f"{', '.join([unit for unit in app.units.keys()])}",
+            parallel=False,
+            coro=app._verify_workload_upgrade(target, list(app.units.values())),
+        ),
+    ]
+    expected_plan.add_steps(upgrade_steps)
+
+    upgrade_plan = app.generate_upgrade_plan(target, False)
+
+    assert_steps(upgrade_plan, expected_plan)
+
+
 def test_auxiliary_upgrade_plan_unknown_track(model):
     """Test auxiliary upgrade plan with unknown track."""
     channel = "2.0/stable"