diff --git a/unit_tests/test_zaza_model.py b/unit_tests/test_zaza_model.py index 0a57bf8..0128b58 100644 --- a/unit_tests/test_zaza_model.py +++ b/unit_tests/test_zaza_model.py @@ -450,7 +450,7 @@ class TestModel(ut_utils.BaseTestCase): 'workload-status-message': 'Unit is ready'}) self.assertTrue( model.check_unit_workload_status(self.Model_mock, - self.unit1, 'active')) + self.unit1, ['active'])) def test_check_unit_workload_status_no_match(self): self.patch_object(model, 'check_model_for_hard_errors') @@ -459,7 +459,17 @@ class TestModel(ut_utils.BaseTestCase): 'workload-status-message': 'Unit is ready'}) self.assertFalse( model.check_unit_workload_status(self.Model_mock, - self.unit1, 'active')) + self.unit1, ['active'])) + + def test_check_unit_workload_status_multi(self): + self.patch_object(model, 'check_model_for_hard_errors') + self._application_states_setup({ + 'workload-status': 'blocked', + 'workload-status-message': 'Unit is ready'}) + self.assertTrue( + model.check_unit_workload_status( + self.Model_mock, + self.unit1, ['active', 'blocked'])) def test_check_unit_workload_status_message_message(self): self.patch_object(model, 'check_model_for_hard_errors') @@ -490,7 +500,7 @@ class TestModel(ut_utils.BaseTestCase): model.check_unit_workload_status_message( self.Model_mock, self.unit1, - prefixes=('Readyish', 'Unit is ready'))) + prefixes=['Readyish', 'Unit is ready'])) def test_check_unit_workload_status_message_prefix_no_match(self): self.patch_object(model, 'check_model_for_hard_errors') @@ -501,7 +511,7 @@ class TestModel(ut_utils.BaseTestCase): model.check_unit_workload_status_message( self.Model_mock, self.unit1, - prefixes=('Readyish', 'Unit is ready'))) + prefixes=['Readyish', 'Unit is ready'])) def test_wait_for_application_states(self): self._application_states_setup({ diff --git a/zaza/model.py b/zaza/model.py index d90b7f7..357c56d 100644 --- a/zaza/model.py +++ b/zaza/model.py @@ -608,7 +608,7 @@ def check_model_for_hard_errors(model): raise UnitError(errored_units) -def check_unit_workload_status(model, unit, state): +def check_unit_workload_status(model, unit, states): """Check that the units workload status matches the supplied state. This function has the side effect of also checking for *any* units @@ -618,14 +618,14 @@ def check_unit_workload_status(model, unit, state): :type model: juju.Model :param unit: Unit to check wl status of :type unit: juju.Unit - :param state: Expected unit work load state - :type state: str + :param states: Acceptable unit work load states + :type states: list :raises: UnitError :returns: Whether units workload status matches desired state :rtype: bool """ check_model_for_hard_errors(model) - return unit.workload_status == state + return unit.workload_status in states def check_unit_workload_status_message(model, unit, message=None, @@ -645,7 +645,7 @@ def check_unit_workload_status_message(model, unit, message=None, :param message: Expected message text :type message: str :param prefixes: Prefixes to match message against - :type prefixes: tuple + :type prefixes: list :raises: ValueError, UnitError :returns: Whether message matches desired string :rtype: bool @@ -654,7 +654,7 @@ def check_unit_workload_status_message(model, unit, message=None, if message is not None: return unit.workload_status_message == message elif prefixes is not None: - return unit.workload_status_message.startswith(prefixes) + return unit.workload_status_message.startswith(tuple(prefixes)) else: raise ValueError("Must be called with message or prefixes") @@ -721,7 +721,8 @@ async def async_wait_for_application_states(model_name=None, states=None, :param timeout: Time to wait for status to be achieved :type timeout: int """ - approved_message_prefixes = ('ready', 'Ready', 'Unit is ready') + approved_message_prefixes = ['ready', 'Ready', 'Unit is ready'] + approved_statuses = ['active'] if not states: states = {} @@ -742,19 +743,25 @@ async def async_wait_for_application_states(model_name=None, states=None, for application, app_data in model.applications.items(): check_info = states.get(application, {}) for unit in app_data.units: + app_wls = check_info.get('workload-status') + if app_wls: + all_approved_statuses = approved_statuses + [app_wls] + else: + all_approved_statuses = approved_statuses logging.info("Checking workload status of {}".format( unit.entity_id)) await model.block_until( lambda: check_unit_workload_status( model, unit, - check_info.get('workload-status', 'active')), + all_approved_statuses), timeout=timeout) check_msg = check_info.get('workload-status-message') logging.info("Checking workload status message of {}" .format(unit.entity_id)) + prefixes = approved_message_prefixes if check_msg is not None: - prefixes = (check_msg) + prefixes = approved_message_prefixes + [check_msg] else: prefixes = approved_message_prefixes await model.block_until(