From 117ae4127d6bb44648388d9be2b948fb747ff683 Mon Sep 17 00:00:00 2001
From: Liam Young <liam.young@canonical.com>
Date: Wed, 28 Mar 2018 15:18:07 +0000
Subject: [PATCH] Add function to scp file to all units

---
 unit_tests/test_zaza_model.py | 14 +++++++++++++
 zaza/model.py                 | 38 ++++++++++++++++++++++++++++++++++-
 2 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/unit_tests/test_zaza_model.py b/unit_tests/test_zaza_model.py
index 010c390..012a55e 100644
--- a/unit_tests/test_zaza_model.py
+++ b/unit_tests/test_zaza_model.py
@@ -9,6 +9,9 @@ class TestModel(ut_utils.BaseTestCase):
 
     def setUp(self):
         super(TestModel, self).setUp()
+
+        async def _scp_to(source, destination, user, proxy, scp_opts):
+            return
         self.unit1 = mock.MagicMock()
         self.unit1.public_address = 'ip1'
         self.unit1.name = 'app/2'
@@ -19,6 +22,8 @@ class TestModel(ut_utils.BaseTestCase):
         self.unit2.name = 'app/4'
         self.unit2.entity_id = 'app/4'
         self.unit2.machine = 'machine7'
+        self.unit1.scp_to.side_effect = _scp_to
+        self.unit2.scp_to.side_effect = _scp_to
         self.units = [self.unit1, self.unit2]
         _units = mock.MagicMock()
         _units.units = self.units
@@ -90,6 +95,15 @@ class TestModel(ut_utils.BaseTestCase):
         unit_mock.scp_to.assert_called_once_with(
             '/tmp/src', '/tmp/dest', proxy=False, scp_opts='', user='ubuntu')
 
+    def test_scp_to_all_units(self):
+        self.patch_object(model, 'Model')
+        self.Model.return_value = self.Model_mock
+        model.scp_to_all_units('app', 'modelname', '/tmp/src', '/tmp/dest')
+        self.unit1.scp_to.assert_called_once_with(
+            '/tmp/src', '/tmp/dest', proxy=False, scp_opts='', user='ubuntu')
+        self.unit2.scp_to.assert_called_once_with(
+            '/tmp/src', '/tmp/dest', proxy=False, scp_opts='', user='ubuntu')
+
     def test_scp_from_unit(self):
         self.patch_object(model, 'Model')
         self.patch_object(model, 'get_unit_from_name')
diff --git a/zaza/model.py b/zaza/model.py
index c5fd41d..db34b12 100644
--- a/zaza/model.py
+++ b/zaza/model.py
@@ -72,7 +72,7 @@ async def run_in_model(model_name, f, add_model_arg=False, awaitable=True):
 
 def scp_to_unit(unit_name, model_name, source, destination, user='ubuntu',
                 proxy=False, scp_opts=''):
-    """Transfer files from to unit_name in model_name.
+    """Transfer files to unit_name in model_name.
 
     :param unit_name: Name of unit to scp to
     :type unit_name: str
@@ -106,6 +106,42 @@ def scp_to_unit(unit_name, model_name, source, destination, user='ubuntu',
         run_in_model(model_name, scp_func, add_model_arg=True, awaitable=True))
 
 
+def scp_to_all_units(application_name, model_name, source, destination,
+                     user='ubuntu', proxy=False, scp_opts=''):
+    """Transfer files from to all units of an application
+
+    :param application_name: Name of application to scp file to
+    :type unit_name: str
+    :param model_name: Name of model unit is in
+    :type model_name: str
+    :param source: Local path of file(s) to transfer
+    :type source: str
+    :param destination: Remote destination of transferred files
+    :type source: str
+    :param user: Remote username
+    :type source: str
+    :param proxy: Proxy through the Juju API server
+    :type proxy: bool
+    :param scp_opts: Additional options to the scp command
+    :type scp_opts: str
+    """
+    async def _scp_to_all_units(application_name, source, destination, user,
+                                proxy, scp_opts, model):
+        for unit in model.applications[application_name].units:
+            await unit.scp_to(source, destination, user=user, proxy=proxy,
+                              scp_opts=scp_opts)
+    scp_func = functools.partial(
+        _scp_to_all_units,
+        application_name,
+        source,
+        destination,
+        user=user,
+        proxy=proxy,
+        scp_opts=scp_opts)
+    loop.run(
+        run_in_model(model_name, scp_func, add_model_arg=True, awaitable=True))
+
+
 def scp_from_unit(unit_name, model_name, source, destination, user='ubuntu',
                   proxy=False, scp_opts=''):
     """Transfer files from to unit_name in model_name.