From 68114544bda8841d7d95484d00c7d9cedde0c8e3 Mon Sep 17 00:00:00 2001
From: Liam Young <liam.young@canonical.com>
Date: Wed, 25 Apr 2018 09:48:31 +0100
Subject: [PATCH] Add the ability to wait for bespoke statuses (#38)

* Add the ability to wait for bespoke statuses

This change adds the ability to wait for bespoke work load
statuses and messages. These are defined in the charms tests.yaml

* Add unit tests

* Remove debug print

* Fix typos and add checks for  errored units

* Restore juju_wait as openstack_utils imports it
---
 .../test_zaza_charm_lifecycle_deploy.py       |  31 +++-
 unit_tests/test_zaza_model.py                 | 171 ++++++++++++++++++
 zaza/charm_lifecycle/README.md                |  13 ++
 zaza/charm_lifecycle/deploy.py                |   8 +-
 zaza/model.py                                 | 154 ++++++++++++++++
 5 files changed, 370 insertions(+), 7 deletions(-)

diff --git a/unit_tests/test_zaza_charm_lifecycle_deploy.py b/unit_tests/test_zaza_charm_lifecycle_deploy.py
index 111be4e..235f6be 100644
--- a/unit_tests/test_zaza_charm_lifecycle_deploy.py
+++ b/unit_tests/test_zaza_charm_lifecycle_deploy.py
@@ -140,6 +140,8 @@ class TestCharmLifecycleDeploy(ut_utils.BaseTestCase):
             ['/tmp/local.yaml'])
 
     def test_deploy_bundle(self):
+        self.patch_object(lc_deploy.utils, 'get_charm_config')
+        self.get_charm_config.return_value = {}
         self.patch_object(lc_deploy, 'render_overlays')
         self.patch_object(lc_deploy.subprocess, 'check_call')
         self.render_overlays.return_value = []
@@ -148,18 +150,39 @@ class TestCharmLifecycleDeploy(ut_utils.BaseTestCase):
             ['juju', 'deploy', '-m', 'newmodel', 'bun.yaml'])
 
     def test_deploy(self):
+        self.patch_object(lc_deploy.zaza.model, 'wait_for_application_states')
+        self.patch_object(lc_deploy.utils, 'get_charm_config')
+        self.get_charm_config.return_value = {}
         self.patch_object(lc_deploy, 'deploy_bundle')
-        self.patch_object(lc_deploy.juju_wait, 'wait')
         lc_deploy.deploy('bun.yaml', 'newmodel')
         self.deploy_bundle.assert_called_once_with('bun.yaml', 'newmodel')
-        self.wait.assert_called_once_with(wait_for_workload=True)
+        self.wait_for_application_states.assert_called_once_with(
+            'newmodel',
+            {})
+
+    def test_deploy_bespoke_states(self):
+        self.patch_object(lc_deploy.zaza.model, 'wait_for_application_states')
+        self.patch_object(lc_deploy.utils, 'get_charm_config')
+        self.get_charm_config.return_value = {
+            'target_deploy_status': {
+                'vault': {
+                    'workload-status': 'blocked',
+                    'workload-status-message': 'Vault needs to be inited'}}}
+        self.patch_object(lc_deploy, 'deploy_bundle')
+        lc_deploy.deploy('bun.yaml', 'newmodel')
+        self.deploy_bundle.assert_called_once_with('bun.yaml', 'newmodel')
+        self.wait_for_application_states.assert_called_once_with(
+            'newmodel',
+            {'vault': {
+                'workload-status': 'blocked',
+                'workload-status-message': 'Vault needs to be inited'}})
 
     def test_deploy_nowait(self):
+        self.patch_object(lc_deploy.zaza.model, 'wait_for_application_states')
         self.patch_object(lc_deploy, 'deploy_bundle')
-        self.patch_object(lc_deploy.juju_wait, 'wait')
         lc_deploy.deploy('bun.yaml', 'newmodel', wait=False)
         self.deploy_bundle.assert_called_once_with('bun.yaml', 'newmodel')
-        self.assertFalse(self.wait.called)
+        self.assertFalse(self.wait_for_application_states.called)
 
     def test_parser(self):
         args = lc_deploy.parse_args([
diff --git a/unit_tests/test_zaza_model.py b/unit_tests/test_zaza_model.py
index 7f8337f..26ee41a 100644
--- a/unit_tests/test_zaza_model.py
+++ b/unit_tests/test_zaza_model.py
@@ -79,9 +79,13 @@ class TestModel(ut_utils.BaseTestCase):
 
         async def _disconnect():
             return
+
         self.Model_mock.connect_model.side_effect = _connect_model
         self.Model_mock.disconnect.side_effect = _disconnect
         self.Model_mock.applications = self.mymodel.applications
+        self.Model_mock.units = {
+            'app/2': self.unit1,
+            'app/4': self.unit2}
 
     def test_run_in_model(self):
         self.patch_object(model, 'Model')
@@ -192,3 +196,170 @@ class TestModel(ut_utils.BaseTestCase):
         self.unit2.run_action.assert_called_once_with(
             'backup',
             backup_dir='/dev/null')
+
+    def _application_states_setup(self, setup, units_idle=True):
+        self.system_ready = True
+
+        async def _block_until(f, timeout=None):
+            result = f()
+            if not result:
+                self.system_ready = False
+            return
+
+        async def _all_units_idle():
+            return units_idle
+        self.Model_mock.block_until.side_effect = _block_until
+        self.patch_object(model, 'Model')
+        self.Model.return_value = self.Model_mock
+        self.Model_mock.all_units_idle.return_value = _all_units_idle
+        p_mock_ws = mock.PropertyMock(
+            return_value=setup['workload-status'])
+        p_mock_wsmsg = mock.PropertyMock(
+            return_value=setup['workload-status-message'])
+        type(self.unit1).workload_status = p_mock_ws
+        type(self.unit1).workload_status_message = p_mock_wsmsg
+        type(self.unit2).workload_status = p_mock_ws
+        type(self.unit2).workload_status_message = p_mock_wsmsg
+
+    def test_units_with_wl_status_state(self):
+        self._application_states_setup({
+            'workload-status': 'active',
+            'workload-status-message': 'Unit is ready'})
+        units = model.units_with_wl_status_state(self.Model_mock, 'active')
+        self.assertTrue(len(units) == 2)
+        self.assertIn(self.unit1, units)
+        self.assertIn(self.unit2, units)
+
+    def test_units_with_wl_status_state_no_match(self):
+        self._application_states_setup({
+            'workload-status': 'blocked',
+            'workload-status-message': 'Unit is ready'})
+        units = model.units_with_wl_status_state(self.Model_mock, 'active')
+        self.assertTrue(len(units) == 0)
+
+    def test_check_model_for_hard_errors(self):
+        self.patch_object(model, 'units_with_wl_status_state')
+        self.units_with_wl_status_state.return_value = []
+        # Test will fail if an Exception is raised
+        model.check_model_for_hard_errors(self.Model_mock)
+
+    def test_check_model_for_hard_errors_found(self):
+        self.patch_object(model, 'units_with_wl_status_state')
+        self.units_with_wl_status_state.return_value = [self.unit1]
+        with self.assertRaises(model.UnitError):
+            model.check_model_for_hard_errors(self.Model_mock)
+
+    def test_check_unit_workload_status(self):
+        self.patch_object(model, 'check_model_for_hard_errors')
+        self._application_states_setup({
+            'workload-status': 'active',
+            'workload-status-message': 'Unit is ready'})
+        self.assertTrue(
+            model.check_unit_workload_status(self.Model_mock,
+                                             self.unit1, 'active'))
+
+    def test_check_unit_workload_status_no_match(self):
+        self.patch_object(model, 'check_model_for_hard_errors')
+        self._application_states_setup({
+            'workload-status': 'blocked',
+            'workload-status-message': 'Unit is ready'})
+        self.assertFalse(
+            model.check_unit_workload_status(self.Model_mock,
+                                             self.unit1, 'active'))
+
+    def test_check_unit_workload_status_message_message(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_message(self.Model_mock,
+                                                     self.unit1,
+                                                     message='Unit is ready'))
+
+    def test_check_unit_workload_status_message_message_not_found(self):
+        self.patch_object(model, 'check_model_for_hard_errors')
+        self._application_states_setup({
+            'workload-status': 'blocked',
+            'workload-status-message': 'Something else'})
+        self.assertFalse(
+            model.check_unit_workload_status_message(self.Model_mock,
+                                                     self.unit1,
+                                                     message='Unit is ready'))
+
+    def test_check_unit_workload_status_message_prefix(self):
+        self.patch_object(model, 'check_model_for_hard_errors')
+        self._application_states_setup({
+            'workload-status': 'blocked',
+            'workload-status-message': 'Unit is ready (OSD Count 23)'})
+        self.assertTrue(
+            model.check_unit_workload_status_message(
+                self.Model_mock,
+                self.unit1,
+                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')
+        self._application_states_setup({
+            'workload-status': 'blocked',
+            'workload-status-message': 'On my holidays'})
+        self.assertFalse(
+            model.check_unit_workload_status_message(
+                self.Model_mock,
+                self.unit1,
+                prefixes=('Readyish', 'Unit is ready')))
+
+    def test_wait_for_application_states(self):
+        self._application_states_setup({
+            'workload-status': 'active',
+            'workload-status-message': 'Unit is ready'})
+        model.wait_for_application_states('modelname', timeout=1)
+        self.assertTrue(self.system_ready)
+
+    def test_wait_for_application_states_not_ready_ws(self):
+        self._application_states_setup({
+            'workload-status': 'blocked',
+            'workload-status-message': 'Unit is ready'})
+        model.wait_for_application_states('modelname', timeout=1)
+        self.assertFalse(self.system_ready)
+
+    def test_wait_for_application_states_not_ready_wsmsg(self):
+        self._application_states_setup({
+            'workload-status': 'active',
+            'workload-status-message': 'Unit is not ready'})
+        model.wait_for_application_states('modelname', timeout=1)
+        self.assertFalse(self.system_ready)
+
+    def test_wait_for_application_states_blocked_ok(self):
+        self._application_states_setup({
+            'workload-status': 'blocked',
+            'workload-status-message': 'Unit is ready'})
+        model.wait_for_application_states(
+            'modelname',
+            states={'app': {
+                'workload-status': 'blocked'}},
+            timeout=1)
+        self.assertTrue(self.system_ready)
+
+    def test_wait_for_application_states_bespoke_msg(self):
+        self._application_states_setup({
+            'workload-status': 'active',
+            'workload-status-message': 'Sure, I could do something'})
+        model.wait_for_application_states(
+            'modelname',
+            states={'app': {
+                'workload-status-message': 'Sure, I could do something'}},
+            timeout=1)
+        self.assertTrue(self.system_ready)
+
+    def test_wait_for_application_states_bespoke_msg_bloked_ok(self):
+        self._application_states_setup({
+            'workload-status': 'blocked',
+            'workload-status-message': 'Sure, I could do something'})
+        model.wait_for_application_states(
+            'modelname',
+            states={'app': {
+                'workload-status': 'blocked',
+                'workload-status-message': 'Sure, I could do something'}},
+            timeout=1)
+        self.assertTrue(self.system_ready)
diff --git a/zaza/charm_lifecycle/README.md b/zaza/charm_lifecycle/README.md
index d96c39c..25514dc 100644
--- a/zaza/charm_lifecycle/README.md
+++ b/zaza/charm_lifecycle/README.md
@@ -174,6 +174,19 @@ smoke_bundles:
   - base-bionic
 ```
 
+ * One of the applications being deployed may have a non-standard workload
+   status target state or message. To inform the deployment step what to
+   wait for an optional target\_deploy\_status stanza can be added:
+
+```
+target_deploy_status:
+  vault:
+    workload-status: blocked
+    workload-status-message: Vault needs to be initialized
+  ntp:
+    workload-status-message: Go for it
+```
+
 # Adding tests to zaza
 
 The setup and tests for a charm should live in zaza, this enables the code to
diff --git a/zaza/charm_lifecycle/deploy.py b/zaza/charm_lifecycle/deploy.py
index ab49bcb..d12d2c1 100755
--- a/zaza/charm_lifecycle/deploy.py
+++ b/zaza/charm_lifecycle/deploy.py
@@ -6,8 +6,7 @@ import subprocess
 import sys
 import tempfile
 
-import juju_wait
-
+import zaza.model
 import zaza.charm_lifecycle.utils as utils
 
 DEFAULT_OVERLAY_TEMPLATE_DIR = 'tests/bundles/overlays'
@@ -222,9 +221,12 @@ def deploy(bundle, model, wait=True):
     """
     deploy_bundle(bundle, model)
     if wait:
+        test_config = utils.get_charm_config()
         logging.info("Waiting for environment to settle")
         utils.set_juju_model(model)
-        juju_wait.wait(wait_for_workload=True)
+        zaza.model.wait_for_application_states(
+            model,
+            test_config.get('target_deploy_status', {}))
 
 
 def parse_args(args):
diff --git a/zaza/model.py b/zaza/model.py
index 9ff7fc8..942a508 100644
--- a/zaza/model.py
+++ b/zaza/model.py
@@ -1,5 +1,6 @@
 import asyncio
 from async_generator import async_generator, yield_, asynccontextmanager
+import logging
 import subprocess
 import yaml
 
@@ -391,6 +392,159 @@ async def async_run_action_on_leader(model_name, application_name, action_name,
 run_action_on_leader = sync_wrapper(async_run_action_on_leader)
 
 
+class UnitError(Exception):
+    """Exception raised for units in error state
+
+    """
+
+    def __init__(self, units):
+        message = "Units {} in error state".format(
+            ','.join([u.entity_id for u in units]))
+        super(UnitError, self).__init__(message)
+
+
+def units_with_wl_status_state(model, state):
+    """Return a list of unit which have a matching workload status
+
+    :returns: Units in error state
+    :rtype: [juju.Unit, ...]
+    """
+    matching_units = []
+    for unit in model.units.values():
+        wl_status = unit.workload_status
+        if wl_status == state:
+            matching_units.append(unit)
+    return matching_units
+
+
+def check_model_for_hard_errors(model):
+    """Check model for any hard errors that should halt a deployment
+
+       The only check currently implemented is checking for units in an
+       error state
+
+    :raises: UnitError
+    """
+    errored_units = units_with_wl_status_state(model, 'error')
+    if errored_units:
+        raise UnitError(errored_units)
+
+
+def check_unit_workload_status(model, unit, state):
+    """Check that the units workload status matches the supplied state.
+       This function has the side effect of also checking for *any* units
+       in an error state and aborting if any are found.
+
+    :param model: Model object to check in
+    :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
+    :raises: UnitError
+    :returns: Whether units workload status matches desired state
+    :rtype: bool
+    """
+    check_model_for_hard_errors(model)
+    return unit.workload_status == state
+
+
+def check_unit_workload_status_message(model, unit, message=None,
+                                       prefixes=None):
+    """Check that the units workload status message matches the supplied
+       message or starts with one of the supplied prefixes. Raises an exception
+       if neither prefixes or message is set. This function has the side effect
+       of also checking for *any* units in an error state and aborting if any
+       are found.
+
+    :param model: Model object to check in
+    :type model: juju.Model
+    :param unit: Unit to check wl status of
+    :type unit: juju.Unit
+    :param message: Expected message text
+    :type message: str
+    :param prefixes: Prefixes to match message against
+    :type prefixes: tuple
+    :raises: ValueError, UnitError
+    :returns: Whether message matches desired string
+    :rtype: bool
+    """
+    check_model_for_hard_errors(model)
+    if message:
+        return unit.workload_status_message == message
+    elif prefixes:
+        return unit.workload_status_message.startswith(prefixes)
+    else:
+        raise ValueError("Must be called with message or prefixes")
+
+
+async def async_wait_for_application_states(model_name, states=None,
+                                            timeout=900):
+    """Wait for model to achieve the desired state
+
+    Check the workload status and workload status message for every unit of
+    every application. By default look for an 'active' workload status and a
+    message that starts with one of the approved_message_prefixes.
+
+    Bespoke statuses and messages can be passed in with states. states takes
+    the form:
+
+    {
+        'app': {
+            'workload-status': 'blocked',
+            'workload-status-message': 'No requests without a prod'}
+        'anotherapp': {
+            'workload-status-message': 'Unit is super ready'}}
+
+
+    :param model_name: Name of model to query.
+    :type model_name: str
+    :param states: Staes to look for
+    :type states: dict
+    :param timeout: Time to wait for status to be achieved
+    :type timeout: int
+    """
+    approved_message_prefixes = ('ready', 'Ready', 'Unit is ready')
+
+    if not states:
+        states = {}
+    async with run_in_model(model_name) as model:
+        check_model_for_hard_errors(model)
+        logging.info("Waiting for all units to be idle")
+        await model.block_until(
+            lambda: model.all_units_idle(), timeout=timeout)
+        for application in model.applications:
+            check_info = states.get(application, {})
+            for unit in model.applications[application].units:
+                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')),
+                    timeout=timeout)
+                check_msg = check_info.get('workload-status-message')
+                logging.info("Checking workload status message of {}".format(
+                    unit.entity_id))
+                if check_msg:
+                    await model.block_until(
+                        lambda: check_unit_workload_status_message(
+                            model,
+                            unit,
+                            message=check_msg),
+                        timeout=timeout)
+                else:
+                    await model.block_until(
+                        lambda: check_unit_workload_status_message(
+                            model,
+                            unit,
+                            prefixes=approved_message_prefixes),
+                        timeout=timeout)
+
+wait_for_application_states = sync_wrapper(async_wait_for_application_states)
+
+
 def get_actions(model_name, application_name):
     """Get the actions an applications supports