From a1f80dced7ea2331e602618a09953a8a1a5b30f4 Mon Sep 17 00:00:00 2001
From: Liam Young <liam.young@canonical.com>
Date: Mon, 26 Mar 2018 08:51:09 +0000
Subject: [PATCH 1/6] Add vault tests

---
 zaza/charm_tests/vault/__init__.py |   0
 zaza/charm_tests/vault/setup.py    |  27 +++++++
 zaza/charm_tests/vault/tests.py    |  90 +++++++++++++++++++++
 zaza/charm_tests/vault/utils.py    | 126 +++++++++++++++++++++++++++++
 4 files changed, 243 insertions(+)
 create mode 100644 zaza/charm_tests/vault/__init__.py
 create mode 100644 zaza/charm_tests/vault/setup.py
 create mode 100644 zaza/charm_tests/vault/tests.py
 create mode 100644 zaza/charm_tests/vault/utils.py

diff --git a/zaza/charm_tests/vault/__init__.py b/zaza/charm_tests/vault/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/zaza/charm_tests/vault/setup.py b/zaza/charm_tests/vault/setup.py
new file mode 100644
index 0000000..5866d65
--- /dev/null
+++ b/zaza/charm_tests/vault/setup.py
@@ -0,0 +1,27 @@
+#!/usr/bin/env python3
+
+import hvac
+import logging
+import os
+import requests
+import time
+import unittest
+import urllib3
+import uuid
+import yaml
+
+import zaza.charm_tests.test_utils as test_utils
+import zaza.charm_tests.vault.utils as vault_utils
+
+def basic_setup():
+    clients = vault_utils.get_clients()
+    unseal_client = clients[0]
+    initialized = vault_utils.is_initialized(unseal_client)
+    # The credentials are written to a file to allow the tests to be re-run
+    # this is mainly useful for manually working on the tests.
+    if initialized:
+       vault_creds = vault_utils.get_credentails()
+    else:
+       vault_creds = vault_utils.init_vault(unseal_client[1])
+       vault_utils.store_credentails(vault_creds)
+    print(vault_creds)
diff --git a/zaza/charm_tests/vault/tests.py b/zaza/charm_tests/vault/tests.py
new file mode 100644
index 0000000..6995013
--- /dev/null
+++ b/zaza/charm_tests/vault/tests.py
@@ -0,0 +1,90 @@
+#!/usr/bin/env python3
+
+import hvac
+import logging
+import os
+import requests
+import time
+import unittest
+import urllib3
+import uuid
+import yaml
+
+import zaza.charm_tests.test_utils as test_utils
+import zaza.charm_tests.vault.utils as vault_utils
+
+
+class VaultTest(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+#        vault_utils = VaultUtils()
+#        cls.clients = vault_utils.get_clients()
+#        unseal_client = cls.clients[0]
+#        initialized = vault_utils.is_initialized(unseal_client)
+#        # The credentials are written to a file to allow the tests to be re-run
+#        # this is mainly useful for manually working on the tests.
+#        auth_file = "/tmp/vault_tests.yaml"
+#        if initialized:
+#            vault_creds = vault_utils.get_credentails_from_file(auth_file)
+#        else:
+#            vault_creds = vault_utils.init_vault(unseal_client[1])
+#            vault_utils.write_credentails(auth_file, vault_creds)
+        cls.clients = vault_utils.get_clients()
+        vault_creds = vault_utils.get_credentails()
+        vault_utils.unseal_all(cls.clients, vault_creds['keys'][0])
+        vault_utils.auth_all(cls.clients, vault_creds['root_token'])
+
+    def test_all_clients_authenticated(self):
+        for (addr, client) in self.clients:
+            for i in range(1, 10):
+                try:
+                    self.assertTrue(client.is_authenticated())
+                except hvac.exceptions.InternalServerError:
+                    time.sleep(2)
+                else:
+                    break
+            else:
+                self.assertTrue(client.is_authenticated())
+
+    def check_read(self, key, value):
+        for (addr, client) in self.clients:
+            self.assertEqual(
+                client.read('secret/uuids')['data']['uuid'],
+                value)
+
+    def test_consistent_read_write(self):
+        key = 'secret/uuids'
+        for (addr, client) in self.clients:
+            value = str(uuid.uuid1())
+            client.write(key, uuid=value, lease='1h')
+            # Now check all clients read the same value back
+            self.check_read(key, value)
+
+    @test_utils.skipIfNotHA('vault')
+    def test_vault_ha_statuses(self):
+        leader = []
+        leader_address = []
+        leader_cluster_address = []
+        for (addr, client) in self.clients:
+            self.assertTrue(client.ha_status['ha_enabled'])
+            leader_address.append(
+                client.ha_status['leader_address'])
+            leader_cluster_address.append(
+                client.ha_status['leader_cluster_address'])
+            if client.ha_status['is_self']:
+                leader.append(addr)
+        # Check there is exactly one leader
+        self.assertEqual(len(leader), 1)
+        # Check both cluster addresses match accross the cluster
+        self.assertEqual(len(set(leader_address)), 1)
+        self.assertEqual(len(set(leader_cluster_address)), 1)
+
+    def test_check_vault_status(self):
+        for (addr, client) in self.clients:
+            self.assertFalse(client.seal_status['sealed'])
+            self.assertTrue(client.seal_status['cluster_name'])
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/zaza/charm_tests/vault/utils.py b/zaza/charm_tests/vault/utils.py
new file mode 100644
index 0000000..eba9739
--- /dev/null
+++ b/zaza/charm_tests/vault/utils.py
@@ -0,0 +1,126 @@
+#!/usr/bin/env python3
+
+import asyncio
+import hvac
+import logging
+import os
+import requests
+import time
+import tempfile
+import unittest
+import urllib3
+import uuid
+import yaml
+
+import zaza.charm_tests.test_utils as test_utils
+import zaza.model
+
+AUTH_FILE = "vault_tests.yaml"
+
+
+def get_client(vault_url):
+    """Return an hvac client for the given URL
+
+    :param vault_url: Vault url to point client at
+    :returns: hvac.Client
+    """
+    return hvac.Client(url=vault_url)
+
+def init_vault(client, shares=1, threshold=1):
+    """Initialise vault
+
+    :param client: hvac.Client Client to use for initiliasation
+    :param shares: int Number of key shares to create
+    :param threshold: int Number of keys needed to unseal vault
+    :returns: hvac.Client
+    """
+    return client.initialize(shares, threshold)
+
+def get_clients(units=None):
+    """Create a list of clients, one per vault server
+
+    :param units: [ip1, ip2, ...] List of IP addresses of vault endpoints
+    :returns: [hvac.Client, ...] List of clients
+    """
+    if not units:
+        units = zaza.model.get_app_ips('vault')
+    clients = []
+    for unit in units:
+        vault_url = 'http://{}:8200'.format(unit)
+        clients.append((unit, get_client(vault_url)))
+    return clients
+
+def is_initialized(client):
+    """Check if vault is initialized
+
+    :param client: hvac.Client Client to use to check if vault is
+                               initialized
+    :returns: bool
+    """
+    initialized = False
+    for i in range(1, 10):
+        try:
+            initialized = client[1].is_initialized()
+        except (ConnectionRefusedError,
+                urllib3.exceptions.NewConnectionError,
+                urllib3.exceptions.MaxRetryError,
+                requests.exceptions.ConnectionError):
+            time.sleep(2)
+        else:
+            break
+    else:
+        raise Exception("Cannot connect")
+    return initialized
+
+def get_credentails():
+    unit = zaza.model.get_first_unit('vault')
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        tmp_file = '{}/{}'.format(tmpdirname, AUTH_FILE)
+        zaza.model.scp_from_unit(unit, '~/{}'.format(AUTH_FILE), tmp_file)
+        with open(tmp_file, 'r') as stream:
+            creds = yaml.load(stream)
+    return creds
+
+def store_credentails(creds):
+    unit = zaza.model.get_first_unit('vault')
+    with tempfile.NamedTemporaryFile(mode='w') as fp:
+        fp.write(yaml.dump(creds))
+        fp.flush()
+        zaza.model.scp_to_unit(unit, fp.name, '~/{}'.format(AUTH_FILE))
+        
+def get_credentails_from_file(auth_file):
+    """Read the vault credentials from the auth_file
+
+    :param auth_file: str Path to file with credentials
+    :returns: {} Dictionary of credentials
+    """
+    with open(auth_file, 'r') as stream:
+        vault_creds = yaml.load(stream)
+    return vault_creds
+
+def write_credentails(auth_file, vault_creds):
+    """Write the vault credentials to the auth_file
+
+    :param auth_file: str Path to file to write credentials
+    """
+    with open(auth_file, 'w') as outfile:
+        yaml.dump(vault_creds, outfile, default_flow_style=False)
+
+def unseal_all(clients, key):
+    """Unseal all the vaults with the given clients with the provided key
+
+    :param clients: [hvac.Client, ...] List of clients
+    :param key: str key to unlock clients
+    """
+    for (addr, client) in clients:
+        if client.is_sealed():
+            client.unseal(key)
+
+def auth_all(clients, token):
+    """Authenticate all the given clients with the provided token
+
+    :param clients: [hvac.Client, ...] List of clients
+    :param token: str token to authorize clients
+    """
+    for (addr, client) in clients:
+        client.token = token

From 8c8bd82c7eafe8f5cfc7e7c93115f44a10244ca7 Mon Sep 17 00:00:00 2001
From: Liam Young <liam.young@canonical.com>
Date: Mon, 26 Mar 2018 09:01:05 +0000
Subject: [PATCH 2/6] Add temp model helpers

---
 zaza/model.py | 60 +++++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 58 insertions(+), 2 deletions(-)

diff --git a/zaza/model.py b/zaza/model.py
index db34b12..9b0f802 100644
--- a/zaza/model.py
+++ b/zaza/model.py
@@ -1,8 +1,12 @@
 import functools
 
 from juju import loop
-from juju.model import Model
+import subprocess
+import logging
+import sys
+import yaml
 
+from juju.model import Model
 
 async def deployed(filter=None):
     # Create a Model instance. We need to connect our Model to a Juju api
@@ -17,6 +21,58 @@ async def deployed(filter=None):
         # Disconnect from the api server and cleanup.
         await model.disconnect()
 
+def get_unit_from_name(unit_name, model):
+    app = unit_name.split('/')[0]
+    for u in model.applications[app].units:
+        if u.entity_id == unit_name:
+            unit = u
+            break
+    else:
+        raise Exception
+    return unit
+
+async def _scp_to_unit(unit_name, source_file, target_file):
+    model = Model()
+    await model.connect_current()
+    try:
+        unit = get_unit_from_name(unit_name, model)
+        #await unit.scp_to(source_file, target_file)
+    finally:
+        # Disconnect from the api server and cleanup.
+        await model.disconnect()
+
+async def _get_first_unit(app_name):
+    model = Model()
+    await model.connect_current()
+    try:
+        unit = model.applications[app_name].units[0]
+    finally:
+        # Disconnect from the api server and cleanup.
+        await model.disconnect()
+    return unit
+
+def scp_from_unit(unit_name, source_file, target_file):
+    cmd = ['juju', 'scp', '{}:{}'.format(unit_name, source_file), target_file]
+    subprocess.check_call(cmd)
+
+def scp_to_unit(unit_name, source_file, target_file):
+    cmd = ['juju', 'scp', source_file, '{}:{}'.format(unit_name, target_file)]
+    subprocess.check_call(cmd)
+
+def get_status():
+    status = subprocess.check_output(['juju', 'status', '--format', 'yaml'])
+    return yaml.load(status)
+
+def get_first_unit(app_name):
+    status = get_status()
+    return sorted(status['applications'][app_name]['units'].keys())[0]
+
+def get_app_ips(app_name):
+    status = get_status()
+    addresses = []
+    for unit in status['applications'][app_name]['units'].values():
+        addresses.append(unit['public-address'])
+    return addresses
 
 def get_unit_from_name(unit_name, model):
     """Return the units that corresponds to the name in the given model
@@ -246,7 +302,7 @@ def get_app_ips(model_name, application_name):
 def main():
     # Run the deploy coroutine in an asyncio event loop, using a helper
     # that abstracts loop creation and teardown.
-    print("Current applications: {}".format(", ".join(loop.run(deployed()))))
+    print("Current applications: {}".format( ", ".join(loop.run(deployed()))))
 
 
 if __name__ == '__main__':

From 79f396eef7f1f498f0cd84f9f1baef6c4022cb92 Mon Sep 17 00:00:00 2001
From: Liam Young <liam.young@canonical.com>
Date: Tue, 27 Mar 2018 15:36:00 +0000
Subject: [PATCH 3/6] Lint fixes

---
 zaza/charm_tests/vault/setup.py |  19 +--
 zaza/charm_tests/vault/tests.py |  17 ---
 zaza/charm_tests/vault/utils.py |  19 +--
 zaza/model.py                   | 239 ++------------------------------
 4 files changed, 25 insertions(+), 269 deletions(-)

diff --git a/zaza/charm_tests/vault/setup.py b/zaza/charm_tests/vault/setup.py
index 5866d65..769c760 100644
--- a/zaza/charm_tests/vault/setup.py
+++ b/zaza/charm_tests/vault/setup.py
@@ -1,18 +1,8 @@
 #!/usr/bin/env python3
 
-import hvac
-import logging
-import os
-import requests
-import time
-import unittest
-import urllib3
-import uuid
-import yaml
-
-import zaza.charm_tests.test_utils as test_utils
 import zaza.charm_tests.vault.utils as vault_utils
 
+
 def basic_setup():
     clients = vault_utils.get_clients()
     unseal_client = clients[0]
@@ -20,8 +10,7 @@ def basic_setup():
     # The credentials are written to a file to allow the tests to be re-run
     # this is mainly useful for manually working on the tests.
     if initialized:
-       vault_creds = vault_utils.get_credentails()
+        vault_creds = vault_utils.get_credentails()
     else:
-       vault_creds = vault_utils.init_vault(unseal_client[1])
-       vault_utils.store_credentails(vault_creds)
-    print(vault_creds)
+        vault_creds = vault_utils.init_vault(unseal_client[1])
+        vault_utils.store_credentails(vault_creds)
diff --git a/zaza/charm_tests/vault/tests.py b/zaza/charm_tests/vault/tests.py
index 6995013..33fae16 100644
--- a/zaza/charm_tests/vault/tests.py
+++ b/zaza/charm_tests/vault/tests.py
@@ -1,14 +1,9 @@
 #!/usr/bin/env python3
 
 import hvac
-import logging
-import os
-import requests
 import time
 import unittest
-import urllib3
 import uuid
-import yaml
 
 import zaza.charm_tests.test_utils as test_utils
 import zaza.charm_tests.vault.utils as vault_utils
@@ -18,18 +13,6 @@ class VaultTest(unittest.TestCase):
 
     @classmethod
     def setUpClass(cls):
-#        vault_utils = VaultUtils()
-#        cls.clients = vault_utils.get_clients()
-#        unseal_client = cls.clients[0]
-#        initialized = vault_utils.is_initialized(unseal_client)
-#        # The credentials are written to a file to allow the tests to be re-run
-#        # this is mainly useful for manually working on the tests.
-#        auth_file = "/tmp/vault_tests.yaml"
-#        if initialized:
-#            vault_creds = vault_utils.get_credentails_from_file(auth_file)
-#        else:
-#            vault_creds = vault_utils.init_vault(unseal_client[1])
-#            vault_utils.write_credentails(auth_file, vault_creds)
         cls.clients = vault_utils.get_clients()
         vault_creds = vault_utils.get_credentails()
         vault_utils.unseal_all(cls.clients, vault_creds['keys'][0])
diff --git a/zaza/charm_tests/vault/utils.py b/zaza/charm_tests/vault/utils.py
index eba9739..8c0859c 100644
--- a/zaza/charm_tests/vault/utils.py
+++ b/zaza/charm_tests/vault/utils.py
@@ -1,18 +1,12 @@
 #!/usr/bin/env python3
 
-import asyncio
 import hvac
-import logging
-import os
 import requests
-import time
 import tempfile
-import unittest
+import time
 import urllib3
-import uuid
 import yaml
 
-import zaza.charm_tests.test_utils as test_utils
 import zaza.model
 
 AUTH_FILE = "vault_tests.yaml"
@@ -26,6 +20,7 @@ def get_client(vault_url):
     """
     return hvac.Client(url=vault_url)
 
+
 def init_vault(client, shares=1, threshold=1):
     """Initialise vault
 
@@ -36,6 +31,7 @@ def init_vault(client, shares=1, threshold=1):
     """
     return client.initialize(shares, threshold)
 
+
 def get_clients(units=None):
     """Create a list of clients, one per vault server
 
@@ -50,6 +46,7 @@ def get_clients(units=None):
         clients.append((unit, get_client(vault_url)))
     return clients
 
+
 def is_initialized(client):
     """Check if vault is initialized
 
@@ -72,6 +69,7 @@ def is_initialized(client):
         raise Exception("Cannot connect")
     return initialized
 
+
 def get_credentails():
     unit = zaza.model.get_first_unit('vault')
     with tempfile.TemporaryDirectory() as tmpdirname:
@@ -81,13 +79,15 @@ def get_credentails():
             creds = yaml.load(stream)
     return creds
 
+
 def store_credentails(creds):
     unit = zaza.model.get_first_unit('vault')
     with tempfile.NamedTemporaryFile(mode='w') as fp:
         fp.write(yaml.dump(creds))
         fp.flush()
         zaza.model.scp_to_unit(unit, fp.name, '~/{}'.format(AUTH_FILE))
-        
+
+
 def get_credentails_from_file(auth_file):
     """Read the vault credentials from the auth_file
 
@@ -98,6 +98,7 @@ def get_credentails_from_file(auth_file):
         vault_creds = yaml.load(stream)
     return vault_creds
 
+
 def write_credentails(auth_file, vault_creds):
     """Write the vault credentials to the auth_file
 
@@ -106,6 +107,7 @@ def write_credentails(auth_file, vault_creds):
     with open(auth_file, 'w') as outfile:
         yaml.dump(vault_creds, outfile, default_flow_style=False)
 
+
 def unseal_all(clients, key):
     """Unseal all the vaults with the given clients with the provided key
 
@@ -116,6 +118,7 @@ def unseal_all(clients, key):
         if client.is_sealed():
             client.unseal(key)
 
+
 def auth_all(clients, token):
     """Authenticate all the given clients with the provided token
 
diff --git a/zaza/model.py b/zaza/model.py
index 9b0f802..9bca14c 100644
--- a/zaza/model.py
+++ b/zaza/model.py
@@ -1,13 +1,10 @@
-import functools
-
 from juju import loop
 import subprocess
-import logging
-import sys
 import yaml
 
 from juju.model import Model
 
+
 async def deployed(filter=None):
     # Create a Model instance. We need to connect our Model to a Juju api
     # server before we can use it.
@@ -21,6 +18,7 @@ async def deployed(filter=None):
         # Disconnect from the api server and cleanup.
         await model.disconnect()
 
+
 def get_unit_from_name(unit_name, model):
     app = unit_name.split('/')[0]
     for u in model.applications[app].units:
@@ -31,6 +29,7 @@ def get_unit_from_name(unit_name, model):
         raise Exception
     return unit
 
+
 async def _scp_to_unit(unit_name, source_file, target_file):
     model = Model()
     await model.connect_current()
@@ -41,6 +40,7 @@ async def _scp_to_unit(unit_name, source_file, target_file):
         # Disconnect from the api server and cleanup.
         await model.disconnect()
 
+
 async def _get_first_unit(app_name):
     model = Model()
     await model.connect_current()
@@ -51,22 +51,27 @@ async def _get_first_unit(app_name):
         await model.disconnect()
     return unit
 
+
 def scp_from_unit(unit_name, source_file, target_file):
     cmd = ['juju', 'scp', '{}:{}'.format(unit_name, source_file), target_file]
     subprocess.check_call(cmd)
 
+
 def scp_to_unit(unit_name, source_file, target_file):
     cmd = ['juju', 'scp', source_file, '{}:{}'.format(unit_name, target_file)]
     subprocess.check_call(cmd)
 
+
 def get_status():
     status = subprocess.check_output(['juju', 'status', '--format', 'yaml'])
     return yaml.load(status)
 
+
 def get_first_unit(app_name):
     status = get_status()
     return sorted(status['applications'][app_name]['units'].keys())[0]
 
+
 def get_app_ips(app_name):
     status = get_status()
     addresses = []
@@ -74,235 +79,11 @@ def get_app_ips(app_name):
         addresses.append(unit['public-address'])
     return addresses
 
-def get_unit_from_name(unit_name, model):
-    """Return the units that corresponds to the name in the given model
-
-    :param unit_name: Name of unit to match
-    :type unit_name: str
-    :param model: Model to perform lookup in
-    :type model: juju.model.Model
-    :returns: Unit matching given name
-    :rtype: juju.unit.Unit or None
-    """
-    app = unit_name.split('/')[0]
-    unit = None
-    for u in model.applications[app].units:
-        if u.entity_id == unit_name:
-            unit = u
-            break
-    else:
-        raise Exception
-    return unit
-
-
-async def run_in_model(model_name, f, add_model_arg=False, awaitable=True):
-    """Run the given function in the model matching the model_name
-
-    :param model_name: Name of model to run function in
-    :type model_name: str
-    :param f: Function to run with given moel in focus
-    :type f: functools.partial
-    :param add_model_arg: Whether to add kwarg pointing at model to the given
-                          function before running it
-    :type add_model_arg: boolean
-    :param awaitable: Whether f is awaitable
-    :type awaitable: boolean
-    :returns: Output of f
-    :rtype: Unknown, depends on the passed in function
-    """
-    model = Model()
-    await model.connect_model(model_name)
-    output = None
-    try:
-        if add_model_arg:
-            f.keywords.update(model=model)
-        if awaitable:
-            output = await f()
-        else:
-            output = f()
-    finally:
-        # Disconnect from the api server and cleanup.
-        await model.disconnect()
-        return output
-
-
-def scp_to_unit(unit_name, model_name, source, destination, user='ubuntu',
-                proxy=False, scp_opts=''):
-    """Transfer files to unit_name in model_name.
-
-    :param unit_name: Name of unit to scp 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_unit(unit_name, source, destination, user, proxy,
-                           scp_opts, model):
-        unit = get_unit_from_name(unit_name, model)
-        await unit.scp_to(source, destination, user=user, proxy=proxy,
-                          scp_opts=scp_opts)
-    scp_func = functools.partial(
-        _scp_to_unit,
-        unit_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_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.
-
-    :param unit_name: Name of unit to scp from
-    :type unit_name: str
-    :param model_name: Name of model unit is in
-    :type model_name: str
-    :param source: Remote path of file(s) to transfer
-    :type source: str
-    :param destination: Local 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_from_unit(unit_name, source, destination, user, proxy,
-                             scp_opts, model):
-        unit = get_unit_from_name(unit_name, model)
-        await unit.scp_from(source, destination, user=user, proxy=proxy,
-                            scp_opts=scp_opts)
-    scp_func = functools.partial(
-        _scp_from_unit,
-        unit_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 get_units(model_name, application_name):
-    """Return all the units of a given application
-
-    :param model_name: Name of model to query.
-    :type model_name: str
-    :param application_name: Name of application to retrieve units for
-    :type application_name: str
-
-    :returns: List of juju units
-    :rtype: [juju.unit.Unit, juju.unit.Unit,...]
-    """
-    async def _get_units(application_name, model):
-        return model.applications[application_name].units
-    f = functools.partial(_get_units, application_name)
-    return loop.run(run_in_model(model_name, f, add_model_arg=True))
-
-
-def get_machines(model_name, application_name):
-    """Return all the machines of a given application
-
-    :param model_name: Name of model to query.
-    :type model_name: str
-    :param application_name: Name of application to retrieve units for
-    :type application_name: str
-
-    :returns: List of juju machines
-    :rtype: [juju.machine.Machine, juju.machine.Machine,...]
-    """
-    async def _get_machines(application_name, model):
-        machines = []
-        for unit in model.applications[application_name].units:
-            machines.append(unit.machine)
-        return machines
-    f = functools.partial(_get_machines, application_name)
-    return loop.run(run_in_model(model_name, f, add_model_arg=True))
-
-
-def get_first_unit_name(model_name, application_name):
-    """Return name of lowest numbered unit of given application
-
-    :param model_name: Name of model to query.
-    :type model_name: str
-    :param application_name: Name of application
-    :type application_name: str
-
-    :returns: Name of lowest numbered unit
-    :rtype: str
-    """
-    return get_units(model_name, application_name)[0].name
-
-
-def get_app_ips(model_name, application_name):
-    """Return public address of all units of an application
-
-    :param model_name: Name of model to query.
-    :type model_name: str
-    :param application_name: Name of application
-    :type application_name: str
-
-    :returns: List of ip addresses
-    :rtype: [str, str,...]
-    """
-    return [u.public_address for u in get_units(model_name, application_name)]
-
 
 def main():
     # Run the deploy coroutine in an asyncio event loop, using a helper
     # that abstracts loop creation and teardown.
-    print("Current applications: {}".format( ", ".join(loop.run(deployed()))))
+    print("Current applications: {}".format(", ".join(loop.run(deployed()))))
 
 
 if __name__ == '__main__':

From fa43b720fbe7ee0459fbe45590b06425f9f87776 Mon Sep 17 00:00:00 2001
From: Liam Young <liam.young@canonical.com>
Date: Wed, 28 Mar 2018 12:03:58 +0000
Subject: [PATCH 4/6] Fix tests given model updates

---
 zaza/charm_tests/vault/utils.py | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/zaza/charm_tests/vault/utils.py b/zaza/charm_tests/vault/utils.py
index 8c0859c..d7ba099 100644
--- a/zaza/charm_tests/vault/utils.py
+++ b/zaza/charm_tests/vault/utils.py
@@ -7,6 +7,7 @@ import time
 import urllib3
 import yaml
 
+import zaza.charm_lifecycle.utils as utils
 import zaza.model
 
 AUTH_FILE = "vault_tests.yaml"
@@ -39,7 +40,7 @@ def get_clients(units=None):
     :returns: [hvac.Client, ...] List of clients
     """
     if not units:
-        units = zaza.model.get_app_ips('vault')
+        units = zaza.model.get_app_ips(utils.get_juju_model(), 'vault')
     clients = []
     for unit in units:
         vault_url = 'http://{}:8200'.format(unit)
@@ -71,21 +72,29 @@ def is_initialized(client):
 
 
 def get_credentails():
-    unit = zaza.model.get_first_unit('vault')
+    unit = zaza.model.get_first_unit_name(utils.get_juju_model(), 'vault')
     with tempfile.TemporaryDirectory() as tmpdirname:
         tmp_file = '{}/{}'.format(tmpdirname, AUTH_FILE)
-        zaza.model.scp_from_unit(unit, '~/{}'.format(AUTH_FILE), tmp_file)
+        zaza.model.scp_from_unit(
+            unit,
+            utils.get_juju_model(),
+            '~/{}'.format(AUTH_FILE),
+            tmp_file)
         with open(tmp_file, 'r') as stream:
             creds = yaml.load(stream)
     return creds
 
 
 def store_credentails(creds):
-    unit = zaza.model.get_first_unit('vault')
+    unit = zaza.model.get_first_unit_name(utils.get_juju_model(), 'vault')
     with tempfile.NamedTemporaryFile(mode='w') as fp:
         fp.write(yaml.dump(creds))
         fp.flush()
-        zaza.model.scp_to_unit(unit, fp.name, '~/{}'.format(AUTH_FILE))
+        zaza.model.scp_to_unit(
+            unit,
+            utils.get_juju_model(),
+            fp.name,
+            '~/{}'.format(AUTH_FILE))
 
 
 def get_credentails_from_file(auth_file):

From 8104b0e7fd4b88d61df35bfe7abc1a16f96a32fb Mon Sep 17 00:00:00 2001
From: Liam Young <liam.young@canonical.com>
Date: Thu, 29 Mar 2018 05:54:13 +0000
Subject: [PATCH 5/6] Pull in new model file

---
 zaza/model.py | 235 ++++++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 199 insertions(+), 36 deletions(-)

diff --git a/zaza/model.py b/zaza/model.py
index 9bca14c..db34b12 100644
--- a/zaza/model.py
+++ b/zaza/model.py
@@ -1,7 +1,6 @@
-from juju import loop
-import subprocess
-import yaml
+import functools
 
+from juju import loop
 from juju.model import Model
 
 
@@ -20,7 +19,17 @@ async def deployed(filter=None):
 
 
 def get_unit_from_name(unit_name, model):
+    """Return the units that corresponds to the name in the given model
+
+    :param unit_name: Name of unit to match
+    :type unit_name: str
+    :param model: Model to perform lookup in
+    :type model: juju.model.Model
+    :returns: Unit matching given name
+    :rtype: juju.unit.Unit or None
+    """
     app = unit_name.split('/')[0]
+    unit = None
     for u in model.applications[app].units:
         if u.entity_id == unit_name:
             unit = u
@@ -30,54 +39,208 @@ def get_unit_from_name(unit_name, model):
     return unit
 
 
-async def _scp_to_unit(unit_name, source_file, target_file):
+async def run_in_model(model_name, f, add_model_arg=False, awaitable=True):
+    """Run the given function in the model matching the model_name
+
+    :param model_name: Name of model to run function in
+    :type model_name: str
+    :param f: Function to run with given moel in focus
+    :type f: functools.partial
+    :param add_model_arg: Whether to add kwarg pointing at model to the given
+                          function before running it
+    :type add_model_arg: boolean
+    :param awaitable: Whether f is awaitable
+    :type awaitable: boolean
+    :returns: Output of f
+    :rtype: Unknown, depends on the passed in function
+    """
     model = Model()
-    await model.connect_current()
+    await model.connect_model(model_name)
+    output = None
     try:
+        if add_model_arg:
+            f.keywords.update(model=model)
+        if awaitable:
+            output = await f()
+        else:
+            output = f()
+    finally:
+        # Disconnect from the api server and cleanup.
+        await model.disconnect()
+        return output
+
+
+def scp_to_unit(unit_name, model_name, source, destination, user='ubuntu',
+                proxy=False, scp_opts=''):
+    """Transfer files to unit_name in model_name.
+
+    :param unit_name: Name of unit to scp 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_unit(unit_name, source, destination, user, proxy,
+                           scp_opts, model):
         unit = get_unit_from_name(unit_name, model)
-        #await unit.scp_to(source_file, target_file)
-    finally:
-        # Disconnect from the api server and cleanup.
-        await model.disconnect()
+        await unit.scp_to(source, destination, user=user, proxy=proxy,
+                          scp_opts=scp_opts)
+    scp_func = functools.partial(
+        _scp_to_unit,
+        unit_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))
 
 
-async def _get_first_unit(app_name):
-    model = Model()
-    await model.connect_current()
-    try:
-        unit = model.applications[app_name].units[0]
-    finally:
-        # Disconnect from the api server and cleanup.
-        await model.disconnect()
-    return unit
+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, source_file, target_file):
-    cmd = ['juju', 'scp', '{}:{}'.format(unit_name, source_file), target_file]
-    subprocess.check_call(cmd)
+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.
+
+    :param unit_name: Name of unit to scp from
+    :type unit_name: str
+    :param model_name: Name of model unit is in
+    :type model_name: str
+    :param source: Remote path of file(s) to transfer
+    :type source: str
+    :param destination: Local 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_from_unit(unit_name, source, destination, user, proxy,
+                             scp_opts, model):
+        unit = get_unit_from_name(unit_name, model)
+        await unit.scp_from(source, destination, user=user, proxy=proxy,
+                            scp_opts=scp_opts)
+    scp_func = functools.partial(
+        _scp_from_unit,
+        unit_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_to_unit(unit_name, source_file, target_file):
-    cmd = ['juju', 'scp', source_file, '{}:{}'.format(unit_name, target_file)]
-    subprocess.check_call(cmd)
+def get_units(model_name, application_name):
+    """Return all the units of a given application
+
+    :param model_name: Name of model to query.
+    :type model_name: str
+    :param application_name: Name of application to retrieve units for
+    :type application_name: str
+
+    :returns: List of juju units
+    :rtype: [juju.unit.Unit, juju.unit.Unit,...]
+    """
+    async def _get_units(application_name, model):
+        return model.applications[application_name].units
+    f = functools.partial(_get_units, application_name)
+    return loop.run(run_in_model(model_name, f, add_model_arg=True))
 
 
-def get_status():
-    status = subprocess.check_output(['juju', 'status', '--format', 'yaml'])
-    return yaml.load(status)
+def get_machines(model_name, application_name):
+    """Return all the machines of a given application
+
+    :param model_name: Name of model to query.
+    :type model_name: str
+    :param application_name: Name of application to retrieve units for
+    :type application_name: str
+
+    :returns: List of juju machines
+    :rtype: [juju.machine.Machine, juju.machine.Machine,...]
+    """
+    async def _get_machines(application_name, model):
+        machines = []
+        for unit in model.applications[application_name].units:
+            machines.append(unit.machine)
+        return machines
+    f = functools.partial(_get_machines, application_name)
+    return loop.run(run_in_model(model_name, f, add_model_arg=True))
 
 
-def get_first_unit(app_name):
-    status = get_status()
-    return sorted(status['applications'][app_name]['units'].keys())[0]
+def get_first_unit_name(model_name, application_name):
+    """Return name of lowest numbered unit of given application
+
+    :param model_name: Name of model to query.
+    :type model_name: str
+    :param application_name: Name of application
+    :type application_name: str
+
+    :returns: Name of lowest numbered unit
+    :rtype: str
+    """
+    return get_units(model_name, application_name)[0].name
 
 
-def get_app_ips(app_name):
-    status = get_status()
-    addresses = []
-    for unit in status['applications'][app_name]['units'].values():
-        addresses.append(unit['public-address'])
-    return addresses
+def get_app_ips(model_name, application_name):
+    """Return public address of all units of an application
+
+    :param model_name: Name of model to query.
+    :type model_name: str
+    :param application_name: Name of application
+    :type application_name: str
+
+    :returns: List of ip addresses
+    :rtype: [str, str,...]
+    """
+    return [u.public_address for u in get_units(model_name, application_name)]
 
 
 def main():

From cb11ddd1e9519ff3567882057188c7726fbb5895 Mon Sep 17 00:00:00 2001
From: Liam Young <liam.young@canonical.com>
Date: Thu, 29 Mar 2018 06:07:12 +0000
Subject: [PATCH 6/6] Fix skip decorator

---
 zaza/charm_tests/test_utils.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/zaza/charm_tests/test_utils.py b/zaza/charm_tests/test_utils.py
index e5e2bd7..81bbf6d 100644
--- a/zaza/charm_tests/test_utils.py
+++ b/zaza/charm_tests/test_utils.py
@@ -1,11 +1,14 @@
 import logging
 import zaza.model
 
+import zaza.charm_lifecycle.utils as utils
+
 
 def skipIfNotHA(service_name):
     def _skipIfNotHA_inner_1(f):
         def _skipIfNotHA_inner_2(*args, **kwargs):
-            if len(zaza.model.get_app_ips(service_name)) > 1:
+            ips = zaza.model.get_app_ips(utils.get_juju_model(), service_name)
+            if len(ips) > 1:
                 return f(*args, **kwargs)
             else:
                 logging.warn("Skipping HA test for non-ha service {}".format(