From 48e627dcac3e44b406fd25aeea7c16b01f682b88 Mon Sep 17 00:00:00 2001
From: Robert Gildein <robert.gildein@canonical.com>
Date: Tue, 12 Mar 2024 13:02:11 +0100
Subject: [PATCH] Add add_steps function to BaseStep

---
 cou/apps/base.py               | 20 +++++++++-----------
 cou/steps/__init__.py          |  9 +++++++++
 tests/unit/steps/test_steps.py | 12 ++++++++++++
 3 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/cou/apps/base.py b/cou/apps/base.py
index 06a883e6..46e470cc 100644
--- a/cou/apps/base.py
+++ b/cou/apps/base.py
@@ -502,11 +502,9 @@ def generate_upgrade_plan(self, target: OpenStackRelease) -> ApplicationUpgradeP
         upgrade_plan = ApplicationUpgradePlan(
             description=f"Upgrade plan for '{self.name}' to {target}",
         )
-        upgrade_plan.sub_steps = [
-            *self.pre_upgrade_steps(target),
-            *self.upgrade_steps(target),
-            *self.post_upgrade_steps(target),
-        ]
+        upgrade_plan.add_steps(self.pre_upgrade_steps(target))
+        upgrade_plan.add_steps(self.upgrade_steps(target))
+        upgrade_plan.add_steps(self.post_upgrade_steps(target))
 
         return upgrade_plan
 
@@ -522,13 +520,13 @@ def _get_upgrade_current_release_packages_step(self) -> PreUpgradeStep:
             ),
             parallel=True,
         )
-        for unit in self.units:
-            step.add_step(
-                UnitUpgradeStep(
-                    description=f"Upgrade software packages on unit {unit.name}",
-                    coro=upgrade_packages(unit.name, self.model, self.packages_to_hold),
-                )
+        step.add_steps(
+            UnitUpgradeStep(
+                description=f"Upgrade software packages on unit {unit.name}",
+                coro=upgrade_packages(unit.name, self.model, self.packages_to_hold),
             )
+            for unit in self.units
+        )
 
         return step
 
diff --git a/cou/steps/__init__.py b/cou/steps/__init__.py
index fa2c7d73..89cdfa10 100644
--- a/cou/steps/__init__.py
+++ b/cou/steps/__init__.py
@@ -231,6 +231,15 @@ def add_step(self, step: BaseStep) -> None:
 
         self._sub_steps.append(step)
 
+    def add_steps(self, steps: Iterable[BaseStep]) -> None:
+        """Add multiple steps.
+
+        :param steps: Sequence of BaseStep to be added as sub steps.
+        :type steps: Iterable[BaseStep]
+        """
+        for step in steps:
+            self.add_step(step)
+
     def cancel(self, safe: bool = True) -> None:
         """Cancel step and all its sub steps.
 
diff --git a/tests/unit/steps/test_steps.py b/tests/unit/steps/test_steps.py
index 93b7b3ea..2b92ecce 100644
--- a/tests/unit/steps/test_steps.py
+++ b/tests/unit/steps/test_steps.py
@@ -228,6 +228,18 @@ def test_step_add_step_failed():
         plan.add_step(MagicMock())
 
 
+def test_step_add_steps():
+    """Test BaseStep adding sub steps at once."""
+    exp_sub_steps = 3
+    plan = BaseStep(description="plan")
+    plan.add_steps(
+        [BaseStep(description=f"sub-step-{i}", coro=mock_coro()) for i in range(exp_sub_steps)]
+        + [BaseStep(description="empty-step")]  # we also check that empty step will not be added
+    )
+
+    assert len(plan.sub_steps) == exp_sub_steps
+
+
 def test_step_cancel_safe():
     """Test step safe cancel."""
     plan = BaseStep(description="plan")