Merge pull request #159 from gnuoy/issue/158

Accept bespoke workload statuses and default

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({

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(