diff --git a/unit_tests/test_zaza_model.py b/unit_tests/test_zaza_model.py index df5c068..eeb1db0 100644 --- a/unit_tests/test_zaza_model.py +++ b/unit_tests/test_zaza_model.py @@ -249,19 +249,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') diff --git a/zaza/model.py b/zaza/model.py index 8656e9d..38c82a7 100644 --- a/zaza/model.py +++ b/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))