Fix enable and disable scheduler order on upgrade

cou/apps/core.py

@@ -19,7 +19,7 @@
 from cou.apps.base import LONG_IDLE_TIMEOUT, OpenStackApplication
 from cou.apps.factory import AppFactory
 from cou.exceptions import ApplicationNotSupported
-from cou.steps import UnitUpgradeStep, UpgradeStep
+from cou.steps import PostUpgradeStep, PreUpgradeStep, UnitUpgradeStep, UpgradeStep
 from cou.utils.juju_utils import Unit
 from cou.utils.nova_compute import verify_empty_hypervisor
 from cou.utils.openstack import OpenStackRelease
@@ -59,6 +59,20 @@ class NovaCompute(OpenStackApplication):
     wait_for_model = True
     upgrade_units_running_vms = False
 
+    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[COUUnit]]
+        :return: List of pre upgrade steps.
+        :rtype: list[PreUpgradeStep]
+        """
+        return self._get_disable_scheduler_step(units) + super().pre_upgrade_steps(target, units)
+
     def upgrade_steps(
         self, target: OpenStackRelease, units: Optional[list[Unit]], force: bool
     ) -> list[UpgradeStep]:
@@ -73,11 +87,27 @@ def upgrade_steps(
         :return: List of upgrade steps.
         :rtype: list[UpgradeStep]
         """
-        if not units:
+        if units is None:
             units = list(self.units.values())
 
         return super().upgrade_steps(target, units, force)
 
+    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]
+        """
+        return self._get_enable_scheduler_step(units) + super().post_upgrade_steps(target, units)
+
     def _get_unit_upgrade_steps(self, unit: Unit, force: bool) -> UnitUpgradeStep:
         """Get the upgrade steps for a single unit.
 
@@ -89,7 +119,6 @@ def _get_unit_upgrade_steps(self, unit: Unit, force: bool) -> UnitUpgradeStep:
         :rtype: UnitUpgradeStep
         """
         unit_plan = UnitUpgradeStep(f"Upgrade plan for unit '{unit.name}'")
-        unit_plan.add_step(self._get_disable_scheduler_step(unit))
 
         if not force:
             unit_plan.add_step(self._get_empty_hypervisor_step(unit))
@@ -98,7 +127,6 @@ def _get_unit_upgrade_steps(self, unit: Unit, force: bool) -> UnitUpgradeStep:
         unit_plan.add_step(self._get_pause_unit_step(unit, is_dependent))
         unit_plan.add_step(self._get_openstack_upgrade_step(unit, is_dependent))
         unit_plan.add_step(self._get_resume_unit_step(unit, is_dependent))
-        unit_plan.add_step(self._get_enable_scheduler_step(unit))
 
         return unit_plan
 
@@ -113,39 +141,47 @@ def _get_empty_hypervisor_step(self, unit: Unit) -> UnitUpgradeStep:
         :rtype: UnitUpgradeStep
         """
         return UnitUpgradeStep(
-            description=f"Verify that unit '{unit.name}' has no VMs running",
+            f"Verify that unit '{unit.name}' has no VMs running",
             coro=verify_empty_hypervisor(unit, self.model),
         )
 
-    def _get_enable_scheduler_step(self, unit: Unit) -> UnitUpgradeStep:
+    def _get_enable_scheduler_step(self, units: Optional[list[Unit]]) -> list[PostUpgradeStep]:
         """Get the step to enable the scheduler, so the unit can create new VMs.
 
-        :param unit: Unit to be enabled.
-        :type unit: Unit
-        :return: Step to enable the scheduler.
-        :rtype: UnitUpgradeStep
+        :param units: Units to be enabled.
+        :type units: Optional[list[Unit]]
+        :return: Steps to enable the scheduler on units
+        :rtype: list[PostUpgradeStep]
         """
-        return UnitUpgradeStep(
-            description=f"Enable nova-compute scheduler from unit: '{unit.name}'",
-            coro=self.model.run_action(
-                unit_name=unit.name, action_name="enable", raise_on_failure=True
-            ),
-        )
-
-    def _get_disable_scheduler_step(self, unit: Unit) -> UnitUpgradeStep:
+        units_to_enable = self.units.values() if units is None else units
+        return [
+            PostUpgradeStep(
+                description=f"Enable nova-compute scheduler from unit: '{unit.name}'",
+                coro=self.model.run_action(
+                    unit_name=unit.name, action_name="enable", raise_on_failure=True
+                ),
+            )
+            for unit in units_to_enable
+        ]
+
+    def _get_disable_scheduler_step(self, units: Optional[list[Unit]]) -> list[PreUpgradeStep]:
         """Get the step to disable the scheduler, so the unit cannot create new VMs.
 
-        :param unit: Unit to be disabled.
-        :type unit: Unit
-        :return: Step to disable the scheduler.
-        :rtype: UnitUpgradeStep
+        :param units: Units to be disabled.
+        :type units:  Optional[list[Unit]]
+        :return: Steps to disable the scheduler on units
+        :rtype: list[PreUpgradeStep]
         """
-        return UnitUpgradeStep(
-            description=f"Disable nova-compute scheduler from unit: '{unit.name}'",
-            coro=self.model.run_action(
-                unit_name=unit.name, action_name="disable", raise_on_failure=True
-            ),
-        )
+        units_to_disable = self.units.values() if units is None else units
+        return [
+            PreUpgradeStep(
+                description=f"Disable nova-compute scheduler from unit: '{unit.name}'",
+                coro=self.model.run_action(
+                    unit_name=unit.name, action_name="disable", raise_on_failure=True
+                ),
+            )
+            for unit in units_to_disable
+        ]
 
 
 @AppFactory.register_application(["swift-proxy", "swift-storage"])

cou/steps/hypervisor.py

@@ -216,6 +216,7 @@ def _generate_upgrade_steps(
 
             units = group.app_units[app.name]
             logger.info("generating upgrade steps for %s units of %s app", app.name, units)
+
             steps.extend(app.upgrade_steps(target, units, force))
 
         return steps
@@ -236,9 +237,7 @@ def _generate_post_upgrade_steps(
         :rtype: list[PostUpgradeStep]
         """
         steps = []
-        # NOTE(rgildein): Using the reverse order of post-upgrade steps, so these steps starts from
-        #                 subordinate or colocated apps and not nova-compute.
-        for app in self.apps[::-1]:
+        for app in self.apps:
             if app.name not in group.app_units:
                 logger.debug(
                     "skipping application %s because it is not part of group %s",

tests/unit/apps/test_core.py

@@ -726,21 +726,47 @@ def test_upgrade_plan_application_already_disable_action_managed(model):
     assert_steps(upgrade_plan, expected_plan)
 
 
+@patch("cou.apps.base.OpenStackApplication._get_refresh_charm_step")
+@patch("cou.apps.base.OpenStackApplication._get_upgrade_current_release_packages_step")
+@patch("cou.apps.core.NovaCompute._get_disable_scheduler_step")
+def test_nova_compute_pre_upgrade_steps(
+    mock_disable, mock_upgrade_package, mock_refresh_charm, model
+):
+    app = _generate_nova_compute_app(model)
+    target = OpenStackRelease("victoria")
+    units = list(app.units.values())
+
+    app.pre_upgrade_steps(target, units)
+    mock_disable.assert_called_once_with(units)
+    mock_upgrade_package.assert_called_once_with(units)
+    mock_refresh_charm.assert_called_once_with(target)
+
+
+@patch("cou.apps.base.OpenStackApplication._get_wait_step")
+@patch("cou.apps.base.OpenStackApplication._get_reached_expected_target_step")
+@patch("cou.apps.core.NovaCompute._get_enable_scheduler_step")
+def test_nova_compute_post_upgrade_steps(mock_enable, mock_expected_target, mock_wait_step, model):
+    app = _generate_nova_compute_app(model)
+    target = OpenStackRelease("victoria")
+    units = list(app.units.values())
+
+    app.post_upgrade_steps(target, units)
+    mock_enable.assert_called_once_with(units)
+    mock_expected_target.assert_called_once_with(target, units)
+    mock_wait_step.assert_called_once_with()
+
+
 @pytest.mark.parametrize("force", [True, False])
 # add_step check if the step added is from BaseStep, so the return is an empty UnitUpgradeStep
-@patch("cou.apps.core.NovaCompute._get_enable_scheduler_step", return_value=UnitUpgradeStep())
 @patch("cou.apps.core.NovaCompute._get_resume_unit_step", return_value=UnitUpgradeStep())
 @patch("cou.apps.core.NovaCompute._get_openstack_upgrade_step", return_value=UnitUpgradeStep())
 @patch("cou.apps.core.NovaCompute._get_pause_unit_step", return_value=UnitUpgradeStep())
 @patch("cou.apps.core.NovaCompute._get_empty_hypervisor_step", return_value=UnitUpgradeStep())
-@patch("cou.apps.core.NovaCompute._get_disable_scheduler_step", return_value=UnitUpgradeStep())
 def test_nova_compute_get_unit_upgrade_steps(
-    mock_disable,
     mock_empty,
     mock_pause,
     mock_upgrade,
     mock_resume,
-    mock_enable,
     model,
     force,
 ):
@@ -749,7 +775,6 @@ def test_nova_compute_get_unit_upgrade_steps(
 
     app._get_unit_upgrade_steps(unit, force)
 
-    mock_disable.assert_called_once_with(unit)
     if force:
         mock_empty.assert_not_called()
     else:
@@ -758,7 +783,6 @@ def test_nova_compute_get_unit_upgrade_steps(
     mock_pause.assert_called_once_with(unit, not force)
     mock_upgrade.assert_called_once_with(unit, not force)
     mock_resume.assert_called_once_with(unit, not force)
-    mock_enable.assert_called_once_with(unit)
 
 
 def test_nova_compute_get_empty_hypervisor_step(model):
@@ -773,28 +797,85 @@ def test_nova_compute_get_empty_hypervisor_step(model):
     assert app._get_empty_hypervisor_step(unit) == expected_step
 
 
-def test_nova_compute_get_enable_scheduler_step(model):
+@pytest.mark.parametrize(
+    "units",
+    [
+        [f"nova-compute/{unit}" for unit in range(1)],
+        [f"nova-compute/{unit}" for unit in range(2)],
+        [f"nova-compute/{unit}" for unit in range(3)],
+    ],
+)
+def test_nova_compute_get_enable_scheduler_step(model, units):
+    """Enable the scheduler on selected units."""
     app = _generate_nova_compute_app(model)
-    units = list(app.units.values())
-    unit = units[0]
+    units_selected = [app.units[unit] for unit in units]
 
-    expected_step = UpgradeStep(
-        description=f"Enable nova-compute scheduler from unit: '{unit.name}'",
-        coro=model.run_action(unit_name=unit.name, action_name="enable", raise_on_failure=True),
-    )
-    assert app._get_enable_scheduler_step(unit) == expected_step
+    expected_step = [
+        PostUpgradeStep(
+            description=f"Enable nova-compute scheduler from unit: '{unit.name}'",
+            coro=model.run_action(
+                unit_name=unit.name, action_name="enable", raise_on_failure=True
+            ),
+        )
+        for unit in units_selected
+    ]
+    assert app._get_enable_scheduler_step(units_selected) == expected_step
 
 
-def test_nova_compute_get_disable_scheduler_step(model):
+def test_nova_compute_get_enable_scheduler_step_no_units(model):
+    """Enable the scheduler on all units if no units are passed."""
     app = _generate_nova_compute_app(model)
-    units = list(app.units.values())
-    unit = units[0]
 
-    expected_step = UpgradeStep(
-        description=f"Disable nova-compute scheduler from unit: '{unit.name}'",
-        coro=model.run_action(unit_name=unit.name, action_name="disable", raise_on_failure=True),
-    )
-    assert app._get_disable_scheduler_step(unit) == expected_step
+    expected_step = [
+        PostUpgradeStep(
+            description=f"Enable nova-compute scheduler from unit: '{unit.name}'",
+            coro=model.run_action(
+                unit_name=unit.name, action_name="enable", raise_on_failure=True
+            ),
+        )
+        for unit in app.units.values()
+    ]
+    assert app._get_enable_scheduler_step(None) == expected_step
+
+
+@pytest.mark.parametrize(
+    "units",
+    [
+        [f"nova-compute/{unit}" for unit in range(1)],
+        [f"nova-compute/{unit}" for unit in range(2)],
+        [f"nova-compute/{unit}" for unit in range(3)],
+    ],
+)
+def test_nova_compute_get_disable_scheduler_step(model, units):
+    """Disable the scheduler on selected units."""
+    app = _generate_nova_compute_app(model)
+    units_selected = [app.units[unit] for unit in units]
+
+    expected_step = [
+        PreUpgradeStep(
+            description=f"Disable nova-compute scheduler from unit: '{unit.name}'",
+            coro=model.run_action(
+                unit_name=unit.name, action_name="disable", raise_on_failure=True
+            ),
+        )
+        for unit in units_selected
+    ]
+    assert app._get_disable_scheduler_step(units_selected) == expected_step
+
+
+def test_nova_compute_get_disable_scheduler_step_no_units(model):
+    """Disable the scheduler on selected units."""
+    app = _generate_nova_compute_app(model)
+    expected_step = [
+        PreUpgradeStep(
+            description=f"Disable nova-compute scheduler from unit: '{unit.name}'",
+            coro=model.run_action(
+                unit_name=unit.name, action_name="disable", raise_on_failure=True
+            ),
+        )
+        for unit in app.units.values()
+    ]
+    assert app._get_disable_scheduler_step(None) == expected_step
 
 
 def _generate_nova_compute_app(model):
@@ -819,6 +900,9 @@ def test_nova_compute_upgrade_plan(model):
     exp_plan = dedent_plan(
         """\
     Upgrade plan for 'nova-compute' to 'victoria'
+        Disable nova-compute scheduler from unit: 'nova-compute/0'
+        Disable nova-compute scheduler from unit: 'nova-compute/1'
+        Disable nova-compute scheduler from unit: 'nova-compute/2'
         Upgrade software packages of 'nova-compute' from the current APT repositories
             Upgrade software packages on unit 'nova-compute/0'
             Upgrade software packages on unit 'nova-compute/1'
@@ -829,26 +913,23 @@ def test_nova_compute_upgrade_plan(model):
         Change charm config of 'nova-compute' 'source' to 'cloud:focal-victoria'
         Upgrade plan for units: nova-compute/0, nova-compute/1, nova-compute/2
             Upgrade plan for unit 'nova-compute/0'
-                Disable nova-compute scheduler from unit: 'nova-compute/0'
                 Verify that unit 'nova-compute/0' has no VMs running
                 ├── Pause the unit: 'nova-compute/0'
                 ├── Upgrade the unit: 'nova-compute/0'
                 ├── Resume the unit: 'nova-compute/0'
-                Enable nova-compute scheduler from unit: 'nova-compute/0'
             Upgrade plan for unit 'nova-compute/1'
-                Disable nova-compute scheduler from unit: 'nova-compute/1'
                 Verify that unit 'nova-compute/1' has no VMs running
                 ├── Pause the unit: 'nova-compute/1'
                 ├── Upgrade the unit: 'nova-compute/1'
                 ├── Resume the unit: 'nova-compute/1'
-                Enable nova-compute scheduler from unit: 'nova-compute/1'
             Upgrade plan for unit 'nova-compute/2'
-                Disable nova-compute scheduler from unit: 'nova-compute/2'
                 Verify that unit 'nova-compute/2' has no VMs running
                 ├── Pause the unit: 'nova-compute/2'
                 ├── Upgrade the unit: 'nova-compute/2'
                 ├── Resume the unit: 'nova-compute/2'
-                Enable nova-compute scheduler from unit: 'nova-compute/2'
+        Enable nova-compute scheduler from unit: 'nova-compute/0'
+        Enable nova-compute scheduler from unit: 'nova-compute/1'
+        Enable nova-compute scheduler from unit: 'nova-compute/2'
         Wait for up to 1800s for model 'test_model' to reach the idle state
         Verify that the workload of 'nova-compute' has been upgraded on units: nova-compute/0, nova-compute/1, nova-compute/2
     """  # noqa: E501 line too long
@@ -888,6 +969,7 @@ def test_nova_compute_upgrade_plan_single_unit(model):
     exp_plan = dedent_plan(
         """\
     Upgrade plan for 'nova-compute' to 'victoria'
+        Disable nova-compute scheduler from unit: 'nova-compute/0'
         Upgrade software packages of 'nova-compute' from the current APT repositories
             Upgrade software packages on unit 'nova-compute/0'
         Refresh 'nova-compute' to the latest revision of 'ussuri/stable'
@@ -896,12 +978,11 @@ def test_nova_compute_upgrade_plan_single_unit(model):
         Change charm config of 'nova-compute' 'source' to 'cloud:focal-victoria'
         Upgrade plan for units: nova-compute/0
             Upgrade plan for unit 'nova-compute/0'
-                Disable nova-compute scheduler from unit: 'nova-compute/0'
                 Verify that unit 'nova-compute/0' has no VMs running
                 ├── Pause the unit: 'nova-compute/0'
                 ├── Upgrade the unit: 'nova-compute/0'
                 ├── Resume the unit: 'nova-compute/0'
-                Enable nova-compute scheduler from unit: 'nova-compute/0'
+        Enable nova-compute scheduler from unit: 'nova-compute/0'
         Wait for up to 1800s for model 'test_model' to reach the idle state
         Verify that the workload of 'nova-compute' has been upgraded on units: nova-compute/0
     """