diff --git a/zaza/charm_lifecycle/test.py b/zaza/charm_lifecycle/test.py
index ab11c77..22bbf85 100644
--- a/zaza/charm_lifecycle/test.py
+++ b/zaza/charm_lifecycle/test.py
@@ -5,6 +5,7 @@ import logging
 import unittest
 import sys
 
+import zaza.model
 import zaza.charm_lifecycle.utils as utils
 
 
@@ -24,7 +25,7 @@ def run_test_list(tests):
 
 def test(model_name, tests):
     """Run all steps to execute tests against the model."""
-    utils.set_juju_model(model_name)
+    zaza.model.set_juju_model(model_name)
     run_test_list(tests)
 
 
diff --git a/zaza/charm_tests/vault/tests.py b/zaza/charm_tests/vault/tests.py
index d9cf522..3b89eca 100644
--- a/zaza/charm_tests/vault/tests.py
+++ b/zaza/charm_tests/vault/tests.py
@@ -61,7 +61,7 @@ class VaultTest(unittest.TestCase):
             '/usr/local/share/ca-certificates/keystone_juju_ca_cert.crt',
             cacert.decode().strip())
         zaza.model.wait_for_application_states(
-            test_config.get('target_deploy_status', {}))
+            states=test_config.get('target_deploy_status', {}))
         ip = zaza.model.get_app_ips(
             'keystone')[0]
         with tempfile.NamedTemporaryFile(mode='w') as fp:
diff --git a/zaza/model.py b/zaza/model.py
index f4fea03..b5ddba2 100644
--- a/zaza/model.py
+++ b/zaza/model.py
@@ -324,7 +324,7 @@ def get_first_unit_name(application_name, model_name=None):
     :returns: Name of lowest numbered unit
     :rtype: str
     """
-    return get_units(model_name, application_name)[0].name
+    return get_units(application_name, model_name=model_name)[0].name
 
 
 def get_app_ips(application_name, model_name=None):
@@ -337,7 +337,8 @@ def get_app_ips(application_name, model_name=None):
     :returns: List of ip addresses
     :rtype: [str, str,...]
     """
-    return [u.public_address for u in get_units(model_name, application_name)]
+    return [u.public_address
+            for u in get_units(application_name, model_name=model_name)]
 
 
 async def async_get_application_config(application_name, model_name=None):