diff --git a/unit_tests/test_zaza_model.py b/unit_tests/test_zaza_model.py index 9771a41..b911888 100644 --- a/unit_tests/test_zaza_model.py +++ b/unit_tests/test_zaza_model.py @@ -112,6 +112,25 @@ class TestModel(ut_utils.BaseTestCase): } self.Model_mock = mock.MagicMock() + # Juju Status Object and data + self.key = "instance-id" + self.key_data = "machine-uuid" + self.machine = "1" + self.machine_data = {self.key: self.key_data} + self.unit = "app/1" + self.unit_data = { + "workload-status": {"status": "active"}, + "machine": self.machine} + self.application = "app" + self.application_data = {"units": {self.unit: self.unit_data}} + self.subordinate_application = "subordinate_application" + self.subordinate_application_data = { + "subordinate-to": [self.application]} + self.juju_status = mock.MagicMock() + self.juju_status.applications = { + self.application: self.application_data} + self.juju_status.machines = self.machine_data + async def _connect_model(model_name): return model_name @@ -886,34 +905,49 @@ disk_formats = ami,ari,aki,vhd,vmdk,raw,qcow2,vdi,iso,root-tar def test_block_until_unit_wl_status(self): async def _block_until(f, timeout=None): - if not f(): + rc = await f() + if not rc: raise asyncio.futures.TimeoutError - self.patch_object(model, 'get_juju_model', return_value='mname') + + async def _get_status(): + return self.juju_status + self.patch_object(model, 'Model') self.Model.return_value = self.Model_mock - self.Model_mock.block_until.side_effect = _block_until + self.patch_object(model, 'get_juju_model', return_value='mname') self.patch_object(model, 'get_unit_from_name') - self.get_unit_from_name.return_value = mock.MagicMock( - workload_status='active') + self.patch_object(model, 'async_get_status') + self.async_get_status.side_effect = _get_status + self.patch_object(model, 'async_block_until') + self.async_block_until.side_effect = _block_until model.block_until_unit_wl_status( - 'app/2', + 'app/1', 'active', timeout=0.1) def test_block_until_unit_wl_status_fail(self): async def _block_until(f, timeout=None): - if not f(): + rc = await f() + if not rc: raise asyncio.futures.TimeoutError - self.patch_object(model, 'get_juju_model', return_value='mname') + + async def _get_status(): + return self.juju_status + + (self.juju_status.applications[self.application] + ["units"][self.unit]["workload-status"]["status"]) = "blocked" + self.patch_object(model, 'Model') self.Model.return_value = self.Model_mock - self.Model_mock.block_until.side_effect = _block_until + self.patch_object(model, 'get_juju_model', return_value='mname') self.patch_object(model, 'get_unit_from_name') - self.get_unit_from_name.return_value = mock.MagicMock( - workload_status='maintenance') + self.patch_object(model, 'async_get_status') + self.async_get_status.side_effect = _get_status + self.patch_object(model, 'async_block_until') + self.async_block_until.side_effect = _block_until with self.assertRaises(asyncio.futures.TimeoutError): model.block_until_unit_wl_status( - 'app/2', + 'app/1', 'active', timeout=0.1) diff --git a/zaza/model.py b/zaza/model.py index 612b702..86e773a 100644 --- a/zaza/model.py +++ b/zaza/model.py @@ -1090,25 +1090,27 @@ async def async_block_until_unit_wl_status(unit_name, status, model_name=None, blocks until the given unit has the desired workload status:: block_until_unit_wl_status( - 'modelname', aunit, - 'active') + 'active' + model_name='modelname') - :param model_name: Name of model to query. - :type model_name: str - :param unit_name: Name of unit to run action on + :param unit_name: Name of unit :type unit_name: str :param status: Status to wait for (active, maintenance etc) :type status: str + :param model_name: Name of model to query. + :type model_name: str :param timeout: Time to wait for unit to achieved desired status :type timeout: float """ - async with run_in_model(model_name) as model: - unit = get_unit_from_name(unit_name, model) - await model.block_until( - lambda: unit.workload_status == status, - timeout=timeout - ) + async def _unit_status(): + app = unit_name.split("/")[0] + model_status = await async_get_status() + return (model_status.applications[app]['units'][unit_name] + ['workload-status']['status'] == status) + + async with run_in_model(model_name): + await async_block_until(_unit_status, timeout=timeout) block_until_unit_wl_status = sync_wrapper( async_block_until_unit_wl_status)