Merge pull request #598 from gnuoy/scale-command-libjuju3.1

Add helper for scaling k8s charms

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')

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.