Make get_unit_from_name use model_name

unit_tests/test_zaza_model.py

@@ -247,19 +247,31 @@ class TestModel(ut_utils.BaseTestCase):
             'app/2')
 
     def test_get_unit_from_name(self):
-        # Normal case
         self.patch_object(model, 'get_juju_model', return_value='mname')
+        self.patch_object(model, 'Model')
+        self.Model.return_value = self.Model_mock
+        # Normal case
+        self.assertEqual(
+            model.get_unit_from_name('app/4', model_name='mname'),
+            self.unit2)
+
+        # Normal case with Model()
         self.assertEqual(
             model.get_unit_from_name('app/4', self.mymodel),
             self.unit2)
 
+        # Normal case, using default
+        self.assertEqual(
+            model.get_unit_from_name('app/4'),
+            self.unit2)
+
         # Unit does not exist
         with self.assertRaises(model.UnitNotFound):
-            model.get_unit_from_name('app/10', self.mymodel)
+            model.get_unit_from_name('app/10', model_name='mname')
 
         # Application does not exist
         with self.assertRaises(model.UnitNotFound):
-            model.get_unit_from_name('bad_name', self.mymodel)
+            model.get_unit_from_name('bad_name', model_name='mname')
 
     def test_get_app_ips(self):
         self.patch_object(model, 'get_juju_model', return_value='mname')

zaza/model.py

@@ -90,20 +90,26 @@ async def deployed():
         await model.disconnect()
 
 
-def get_unit_from_name(unit_name, model):
+def get_unit_from_name(unit_name, model=None, model_name=None):
     """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
+    :type model: model.Model()
+    :param model_name: Name of the model to perform lookup in
+    :type model_name: string
     :returns: Unit matching given name
     :rtype: juju.unit.Unit or None
+    :raises: UnitNotFound
     """
     app = unit_name.split('/')[0]
     unit = None
     try:
-        units = model.applications[app].units
+        if model is not None:
+            units = model.applications[app].units
+        else:
+            units = get_units(application_name=app, model_name=model_name)
     except KeyError:
         msg = ('Application: {} does not exist in current model'.
                format(app))