From 9747fc257b19e96f6c3f05ea5086be7c1542a2bc Mon Sep 17 00:00:00 2001
From: Liam Young <liam.young@canonical.com>
Date: Mon, 6 Mar 2023 14:57:32 +0000
Subject: [PATCH] Add helper for scaling k8s charms

Add a wrapper to the scale command so k8s applications can be
scaled in and out.

(cherry picked from commit 05cd81d1e0915bf82a32fe678c24329c6a72a5e9)
---
 unit_tests/test_zaza_model.py | 34 ++++++++++++++++++++++++++++++++++
 zaza/model.py                 | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 68 insertions(+)

diff --git a/unit_tests/test_zaza_model.py b/unit_tests/test_zaza_model.py
index cc1ae1cac..8fa8f3316 100644
--- a/unit_tests/test_zaza_model.py
+++ b/unit_tests/test_zaza_model.py
@@ -134,6 +134,9 @@ async def _add_unit(count=1, to=None):
         async def _destroy_unit(*unitnames):
             return
 
+        async def _scale(scale=None, scale_change=None):
+            return
+
         def _is_leader(leader):
             async def _inner_is_leader():
                 return leader
@@ -213,6 +216,7 @@ def fail_on_use():
         _units.destroy_relation.side_effect = _destroy_relation
         _units.add_unit.side_effect = _add_unit
         _units.destroy_unit.side_effect = _destroy_unit
+        _units.scale.side_effect = _scale
 
         self.mymodel = mock.MagicMock()
         self.mymodel.applications = {
@@ -974,6 +978,36 @@ def test_destroy_unit_wait(self):
         self.async_block_until_unit_count.assert_called_once_with(
             'app', 1, model_name=None)
 
+    def test_scale_out(self):
+        self.patch_object(model, 'get_juju_model', return_value='mname')
+        self.patch_object(model, 'Model')
+        self.Model.return_value = self.Model_mock
+        model.scale('app', scale=3)
+        self.mymodel.applications['app'].scale.assert_called_once_with(
+            scale=3, scale_change=None)
+
+    def test_scale_wait(self):
+        self.patch_object(model, 'async_block_until_unit_count')
+        self.patch_object(model, 'get_juju_model', return_value='mname')
+        self.patch_object(model, 'Model')
+        self.Model.return_value = self.Model_mock
+        model.scale('app', scale=3, wait=True)
+        self.mymodel.applications['app'].scale.assert_called_once_with(
+            scale=3, scale_change=None)
+        self.async_block_until_unit_count.assert_called_once_with(
+            'app', 3, model_name=None)
+
+    def test_scale_back_wait(self):
+        self.patch_object(model, 'async_block_until_unit_count')
+        self.patch_object(model, 'get_juju_model', return_value='mname')
+        self.patch_object(model, 'Model')
+        self.Model.return_value = self.Model_mock
+        model.scale('app', scale_change=-1, wait=True)
+        self.mymodel.applications['app'].scale.assert_called_once_with(
+            scale=None, scale_change=-1)
+        self.async_block_until_unit_count.assert_called_once_with(
+            'app', 1, model_name=None)
+
     def test_get_relation_id_interface(self):
         self.patch_object(model, 'get_juju_model', return_value='mname')
         self.patch_object(model, 'Model')
diff --git a/zaza/model.py b/zaza/model.py
index 0d33b7931..01048b3d6 100644
--- a/zaza/model.py
+++ b/zaza/model.py
@@ -2780,6 +2780,40 @@ async def async_destroy_unit(application_name, *unit_names, model_name=None,
 destroy_unit = sync_wrapper(async_destroy_unit)
 
 
+async def async_scale(application_name, scale=None, scale_change=None,
+                      wait=False, model_name=None):
+    """
+    Set or adjust the scale of this (K8s) application.
+
+    :param application_name: Name of application to add unit(s) to
+    :type application_name: str
+    :param scale: Scale to which to set this application.
+    :type scale: int
+    :param scale_change: Amount by which to adjust the scale of this
+            application (can be positive or negative).
+    :type scale_change: int
+    :param wait: Whether to wait for the unit change to appear in juju
+                 status
+    :type wait: bool
+    :param model_name: Name of model to operate on.
+    :type model_name: str
+    """
+    model = await get_model(model_name)
+    app = model.applications[application_name]
+    await app.scale(scale=scale, scale_change=scale_change)
+    if wait:
+        if scale:
+            target_count = scale
+        else:
+            target_count = len(app.units) + scale_change
+        await async_block_until_unit_count(
+            application_name,
+            target_count,
+            model_name=model_name)
+
+scale = sync_wrapper(async_scale)
+
+
 def set_model_constraints(constraints, model_name=None):
     """
     Set constraints on a model.