add UnitUpgradeStep, improve add_step

cou/steps/__init__.py

@@ -21,7 +21,7 @@
 import logging
 import os
 import warnings
-from typing import Any, Coroutine, List, Optional
+from typing import Any, Coroutine, Iterable, List, Optional
 
 from cou.exceptions import CanceledStep
 
@@ -90,7 +90,7 @@ def __init__(
         self._coro: Optional[Coroutine] = coro
         self.parallel = parallel
         self.description = description
-        self.sub_steps: List[BaseStep] = []
+        self._sub_steps: List[BaseStep] = []
         self._canceled: bool = False
         self._task: Optional[asyncio.Task] = None
 
@@ -195,13 +195,36 @@ def done(self) -> bool:
 
         return self._task.done()
 
+    @property
+    def sub_steps(self) -> List[BaseStep]:
+        """Return list of sub-steps.
+
+        :return: List of BaseStep.
+        :rtype: List[BaseStep]
+        """
+        return self._sub_steps
+
+    @sub_steps.setter
+    def sub_steps(self, steps: Iterable[BaseStep]) -> None:
+        """Set a list of sub-steps.
+
+        :param steps: Iterable object containing all steps.
+        :type steps: Iterable
+        """
+        for step in steps:
+            self.add_step(step)
+
     def add_step(self, step: BaseStep) -> None:
         """Add a single step.
 
         :param step: BaseStep to be added as sub step.
         :type step: BaseStep
+        :raises TypeError: If step is not based on BaseStep.
         """
-        self.sub_steps.append(step)
+        if not isinstance(step, BaseStep):
+            raise TypeError("Cannot add an upgrade step that is not derived from BaseStep")
+
+        self._sub_steps.append(step)
 
     def cancel(self, safe: bool = True) -> None:
         """Cancel step and all its sub steps.
@@ -286,9 +309,24 @@ class UpgradeStep(BaseStep):
     prompt: bool = False
 
 
+class UnitUpgradeStep(UpgradeStep):
+    """Represents the upgrade step for an individual unit."""
+
+
 class PreUpgradeStep(UpgradeStep):
     """Represents the pre-upgrade step."""
 
 
 class PostUpgradeStep(UpgradeStep):
     """Represents the post-upgrade step."""
+
+
+def is_unit_upgrade_step(step: BaseStep) -> bool:
+    """Check if it's unit upgrade step.
+
+    :param step: step
+    :type step: BaseStep
+    :return: True if it's UnitUpgradeStep
+    :rtype: bool
+    """
+    return isinstance(step, UnitUpgradeStep)

tests/unit/steps/test_steps.py

@@ -24,8 +24,11 @@
     BaseStep,
     PostUpgradeStep,
     PreUpgradeStep,
+    UnitUpgradeStep,
     UpgradePlan,
+    UpgradeStep,
     compare_step_coroutines,
+    is_unit_upgrade_step,
 )
 
 
@@ -208,6 +211,15 @@ def test_step_add_step():
     assert len(plan.sub_steps) == exp_sub_steps
 
 
+def test_step_add_step_failed():
+    """Test BaseStep adding sub steps failing."""
+    exp_error_msg = "only steps that are derived from BaseStep are supported"
+    plan = BaseStep(description="plan")
+
+    with pytest.raises(TypeError, match=exp_error_msg):
+        plan.add_step(MagicMock())
+
+
 def test_step_cancel_safe():
     """Test step safe cancel."""
     plan = BaseStep(description="plan")
@@ -367,3 +379,17 @@ async def step_run(_step):
         await step_run(plan)
 
     assert steps_order == exp_order
+
+
+@pytest.mark.parametrize(
+    "step, exp_result",
+    [
+        (UnitUpgradeStep(), True),
+        (UpgradeStep(), False),
+        (PreUpgradeStep(), False),
+        (BaseStep(), False),
+    ],
+)
+def test_is_unit_upgrade_step(step, exp_result):
+    """Test helper function for checking if step is unit upgrade step."""
+    assert is_unit_upgrade_step(step) == exp_result

tests/unit/steps/test_steps_plan.py

@@ -252,7 +252,8 @@ def test_determine_upgrade_target_release_out_of_range(current_os_release, curre
 @pytest.mark.asyncio
 async def test_create_upgrade_plan():
     """Test create_upgrade_group."""
-    app: OpenStackApplication = MagicMock(spec=OpenStackApplication)
+    app: OpenStackApplication = MagicMock(spec_set=OpenStackApplication)
+    app.generate_upgrade_plan.return_value = MagicMock(spec_set=ApplicationUpgradePlan)
     target = OpenStackRelease("victoria")
     description = "test"