From cf1ea4c71b9bcc1bec13362f33f28dbf768133a5 Mon Sep 17 00:00:00 2001
From: Chris MacNaughton <chris.macnaughton@canonical.com>
Date: Wed, 8 Apr 2020 16:10:38 +0200
Subject: [PATCH 01/19] First pass at batched parallel series upgrade

---
 .../test_zaza_utilities_upgrade_utils.py      |   4 +-
 .../series_upgrade/parallel_tests.py          | 202 +++++++
 .../utilities/parallel_series_upgrade.py      | 528 ++++++++++++++++++
 zaza/openstack/utilities/upgrade_utils.py     |   6 +-
 4 files changed, 737 insertions(+), 3 deletions(-)
 create mode 100644 zaza/openstack/charm_tests/series_upgrade/parallel_tests.py
 create mode 100755 zaza/openstack/utilities/parallel_series_upgrade.py

diff --git a/unit_tests/utilities/test_zaza_utilities_upgrade_utils.py b/unit_tests/utilities/test_zaza_utilities_upgrade_utils.py
index 1f21b49..814772e 100644
--- a/unit_tests/utilities/test_zaza_utilities_upgrade_utils.py
+++ b/unit_tests/utilities/test_zaza_utilities_upgrade_utils.py
@@ -90,6 +90,7 @@ class TestUpgradeUtils(ut_utils.BaseTestCase):
 
     def test_get_upgrade_groups(self):
         expected = collections.OrderedDict([
+            ('Stateful Services', []),
             ('Core Identity', []),
             ('Control Plane', ['cinder']),
             ('Data Plane', ['nova-compute']),
@@ -103,10 +104,11 @@ class TestUpgradeUtils(ut_utils.BaseTestCase):
 
     def test_get_series_upgrade_groups(self):
         expected = collections.OrderedDict([
+            ('Stateful Services', ['mydb']),
             ('Core Identity', []),
             ('Control Plane', ['cinder']),
             ('Data Plane', ['nova-compute']),
-            ('sweep_up', ['mydb', 'ntp'])])
+            ('sweep_up', ['ntp'])])
         actual = openstack_upgrade.get_series_upgrade_groups()
         pprint.pprint(expected)
         pprint.pprint(actual)
diff --git a/zaza/openstack/charm_tests/series_upgrade/parallel_tests.py b/zaza/openstack/charm_tests/series_upgrade/parallel_tests.py
new file mode 100644
index 0000000..1b64994
--- /dev/null
+++ b/zaza/openstack/charm_tests/series_upgrade/parallel_tests.py
@@ -0,0 +1,202 @@
+#!/usr/bin/env python3
+
+# Copyright 2018 Canonical Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Define class for Series Upgrade."""
+
+import asyncio
+import logging
+import os
+import sys
+import unittest
+
+from zaza import model
+from zaza.openstack.utilities import (
+    cli as cli_utils,
+    upgrade_utils as upgrade_utils,
+)
+from zaza.openstack.charm_tests.nova.tests import LTSGuestCreateTest
+from zaza.openstack.utilities.parallel_series_upgrade import (
+    parallel_series_upgrade,
+)
+
+
+def _filter_easyrsa(app, app_config, model_name=None):
+    charm_name = upgrade_utils.extract_charm_name_from_url(app_config['charm'])
+    if "easyrsa" in charm_name:
+        logging.warn("Skipping series upgrade of easyrsa Bug #1850121")
+        return True
+    return False
+
+
+def _filter_etcd(app, app_config, model_name=None):
+    charm_name = upgrade_utils.extract_charm_name_from_url(app_config['charm'])
+    if "etcd" in charm_name:
+        logging.warn("Skipping series upgrade of easyrsa Bug #1850124")
+        return True
+    return False
+
+
+class ParallelSeriesUpgradeTest(unittest.TestCase):
+    """Class to encapsulate Series Upgrade Tests."""
+
+    @classmethod
+    def setUpClass(cls):
+        """Run setup for Series Upgrades."""
+        cli_utils.setup_logging()
+        cls.from_series = None
+        cls.to_series = None
+        cls.workaround_script = None
+        cls.files = []
+
+    def test_200_run_series_upgrade(self):
+        """Run series upgrade."""
+        # Set Feature Flag
+        os.environ["JUJU_DEV_FEATURE_FLAGS"] = "upgrade-series"
+        upgrade_groups = upgrade_utils.get_series_upgrade_groups(
+            extra_filters=[_filter_etcd, _filter_easyrsa])
+        from_series = self.from_series
+        to_series = self.to_series
+        completed_machines = []
+        workaround_script = None
+        files = []
+        applications = model.get_status().applications
+        for group_name, apps in upgrade_groups.items():
+            logging.info("About to upgrade {} from {} to {}".format(
+                group_name, from_series, to_series))
+            upgrade_functions = []
+            if group_name in ["Stateful Services", "Data Plane", "sweep_up"]:
+                logging.info("Going to upgrade {} unit by unit".format(apps))
+                upgrade_function = \
+                    parallel_series_upgrade.serial_series_upgrade
+            else:
+                logging.info("Going to upgrade {} all at once".format(apps))
+                upgrade_function = \
+                    parallel_series_upgrade.parallel_series_upgrade
+
+            for charm_name in apps:
+                charm = applications[charm_name]['charm']
+                name = upgrade_utils.extract_charm_name_from_url(charm)
+                upgrade_config = parallel_series_upgrade.app_config(name)
+                upgrade_functions.append(
+                    upgrade_function(
+                        charm_name,
+                        **upgrade_config,
+                        from_series=from_series,
+                        to_series=to_series,
+                        completed_machines=completed_machines,
+                        workaround_script=workaround_script,
+                        files=files))
+            asyncio.get_event_loop().run_until_complete(
+                asyncio.gather(*upgrade_functions))
+
+            if "rabbitmq-server" in apps:
+                logging.info(
+                    "Running complete-cluster-series-upgrade action on leader")
+                model.run_action_on_leader(
+                    'rabbitmq-server',
+                    'complete-cluster-series-upgrade',
+                    action_params={})
+
+            if "percona-cluster" in apps:
+                logging.info(
+                    "Running complete-cluster-series-upgrade action on leader")
+                model.run_action_on_leader(
+                    'mysql',
+                    'complete-cluster-series-upgrade',
+                    action_params={})
+            model.block_until_all_units_idle()
+            logging.info("Finished {}".format(group_name))
+        logging.info("Done!")
+
+
+class OpenStackParallelSeriesUpgrade(ParallelSeriesUpgradeTest):
+    """OpenStack Series Upgrade.
+
+    Full OpenStack series upgrade with VM launch before and after the series
+    upgrade.
+
+    This test requires a full OpenStack including at least: keystone, glance,
+    nova-cloud-controller, nova-compute, neutron-gateway, neutron-api and
+    neutron-openvswitch.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        """Run setup for Series Upgrades."""
+        super(OpenStackParallelSeriesUpgrade, cls).setUpClass()
+        cls.lts = LTSGuestCreateTest()
+        cls.lts.setUpClass()
+
+    def test_100_validate_pre_series_upgrade_cloud(self):
+        """Validate pre series upgrade."""
+        logging.info("Validate pre-series-upgrade: Spin up LTS instance")
+        self.lts.test_launch_small_instance()
+
+    def test_500_validate_series_upgraded_cloud(self):
+        """Validate post series upgrade."""
+        logging.info("Validate post-series-upgrade: Spin up LTS instance")
+        self.lts.test_launch_small_instance()
+
+
+class TrustyXenialSeriesUpgrade(OpenStackParallelSeriesUpgrade):
+    """OpenStack Trusty to Xenial Series Upgrade."""
+
+    @classmethod
+    def setUpClass(cls):
+        """Run setup for Trusty to Xenial Series Upgrades."""
+        super(TrustyXenialSeriesUpgrade, cls).setUpClass()
+        cls.from_series = "trusty"
+        cls.to_series = "xenial"
+
+
+class XenialBionicSeriesUpgrade(OpenStackParallelSeriesUpgrade):
+    """OpenStack Xenial to Bionic Series Upgrade."""
+
+    @classmethod
+    def setUpClass(cls):
+        """Run setup for Xenial to Bionic Series Upgrades."""
+        super(XenialBionicSeriesUpgrade, cls).setUpClass()
+        cls.from_series = "xenial"
+        cls.to_series = "bionic"
+
+
+class BionicFocalSeriesUpgrade(OpenStackParallelSeriesUpgrade):
+    """OpenStack Bionic to FocalSeries Upgrade."""
+
+    @classmethod
+    def setUpClass(cls):
+        """Run setup for Xenial to Bionic Series Upgrades."""
+        super(BionicFocalSeriesUpgrade, cls).setUpClass()
+        cls.from_series = "bionic"
+        cls.to_series = "focal"
+
+
+if __name__ == "__main__":
+    from_series = os.environ.get("FROM_SERIES")
+    if from_series == "trusty":
+        to_series = "xenial"
+        series_upgrade_test = TrustyXenialSeriesUpgrade()
+    elif from_series == "xenial":
+        to_series = "bionic"
+        series_upgrade_test = XenialBionicSeriesUpgrade()
+    elif from_series == "bionic":
+        to_series = "focal"
+        series_upgrade_test = BionicFocalSeriesUpgrade()
+
+    else:
+        raise Exception("FROM_SERIES is not set to a vailid LTS series")
+    series_upgrade_test.setUpClass()
+    sys.exit(series_upgrade_test.test_200_run_series_upgrade())
diff --git a/zaza/openstack/utilities/parallel_series_upgrade.py b/zaza/openstack/utilities/parallel_series_upgrade.py
new file mode 100755
index 0000000..a959838
--- /dev/null
+++ b/zaza/openstack/utilities/parallel_series_upgrade.py
@@ -0,0 +1,528 @@
+#!/usr/bin/env python3
+# Copyright 2020 Canonical Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Collection of functions for testing series upgrade in parallel."""
+
+
+import asyncio
+import collections
+import copy
+import logging
+import subprocess
+
+from zaza import model
+import zaza.openstack.utilities.generic as os_utils
+import zaza.openstack.utilities.series_upgrade as series_upgrade_utils
+from zaza.openstack.utilities.series_upgrade import (
+    SUBORDINATE_PAUSE_RESUME_BLACKLIST,
+)
+
+
+def app_config(charm_name):
+    """Return a dict with the upgrade config for an application.
+
+    :param charm_name: Name of the charm about to upgrade
+    :type charm_name: str
+    :param async: Whether the upgreade functions should be async
+    :type async: bool
+    :returns: A dicitonary of the upgrade config for the application
+    :rtype: Dict
+    """
+    default = {
+        'origin': 'openstack-origin',
+        'pause_non_leader_subordinate': True,
+        'pause_non_leader_primary': True,
+        'post_upgrade_functions': [],
+        'follower_first': False, }
+    _app_settings = collections.defaultdict(lambda: default)
+    ceph = {
+        'origin': "source",
+        'pause_non_leader_primary': False,
+        'pause_non_leader_subordinate': False,
+    }
+    exceptions = {
+        'rabbitmq-server': {
+            'origin': 'source',
+            'pause_non_leader_subordinate': False, },
+        'percona-cluster': {'origin': 'source', },
+        'nova-compute': {
+            'pause_non_leader_primary': False,
+            'pause_non_leader_subordinate': False, },
+        'ceph': ceph,
+        'ceph-mon': ceph,
+        'ceph-osd': ceph,
+        'designate-bind': {'origin': None, },
+        'tempest': {'origin': None, },
+        'memcached': {
+            'origin': None,
+            'pause_non_leader_primary': False,
+            'pause_non_leader_subordinate': False,
+        },
+        'vault': {
+            'origin': None,
+            'pause_non_leader_primary': False,
+            'pause_non_leader_subordinate': True,
+            'post_upgrade_functions': [
+                ('zaza.openstack.charm_tests.vault.setup.'
+                 'mojo_unseal_by_unit')]
+        },
+        'mongodb': {
+            'origin': None,
+            'follower_first': True,
+        }
+    }
+    for key, value in exceptions.items():
+        _app_settings[key] = copy.deepcopy(default)
+        _app_settings[key].update(value)
+    return _app_settings[charm_name]
+
+
+def upgrade_ubuntu_lite(from_series='xenial', to_series='bionic'):
+    """Validate that we can upgrade the ubuntu-lite charm.
+
+    :param from_series: What series are we upgrading from
+    :type from_series: str
+    :param to_series: What series are we upgrading to
+    :type to_series: str
+    """
+    completed_machines = []
+    asyncio.get_event_loop().run_until_complete(
+        parallel_series_upgrade(
+            'ubuntu-lite', pause_non_leader_primary=False,
+            pause_non_leader_subordinate=False,
+            completed_machines=completed_machines, origin=None)
+    )
+
+
+async def parallel_series_upgrade(
+        application, from_series='xenial', to_series='bionic',
+        origin='openstack-origin', pause_non_leader_primary=True,
+        pause_non_leader_subordinate=True, post_upgrade_functions=None,
+        completed_machines=None, files=None, workaround_script=None,
+        follower_first=False):
+    """Perform series upgrade on an application in parallel.
+
+    :param unit_name: Unit Name
+    :type unit_name: str
+    :param machine_num: Machine number
+    :type machine_num: str
+    :param from_series: The series from which to upgrade
+    :type from_series: str
+    :param to_series: The series to which to upgrade
+    :type to_series: str
+    :param origin: The configuration setting variable name for changing origin
+                   source. (openstack-origin or source)
+    :type origin: str
+    :param files: Workaround files to scp to unit under upgrade
+    :type files: list
+    :param workaround_script: Workaround script to run during series upgrade
+    :type workaround_script: str
+    :param follower_first: Should the follower(s) be upgraded first
+    :type follower_first: bool
+    :returns: None
+    :rtype: None
+    """
+    if completed_machines is None:
+        completed_machines = []
+    if files is None:
+        files = []
+    if post_upgrade_functions is None:
+        post_upgrade_functions = []
+    if follower_first:
+        logging.error("leader_first is ignored for parallel upgrade")
+    logging.info(
+        "About to upgrade the units of {} in parallel (follower first: {})"
+        .format(application, follower_first))
+    # return
+    status = (await model.async_get_status()).applications[application]
+    leader, non_leaders = await get_leader_and_non_leaders(application)
+    for leader_name, leader_unit in leader.items():
+        leader_machine = leader_unit["machine"]
+        leader = leader_name
+    machines = [
+        unit["machine"] for name, unit
+        in non_leaders.items()
+        if unit['machine'] not in completed_machines]
+    await maybe_pause_things(
+        status,
+        non_leaders,
+        pause_non_leader_subordinate,
+        pause_non_leader_primary)
+    await series_upgrade_utils.async_set_series(
+        application, to_series=to_series)
+
+    prepare_group = [
+        prepare_series_upgrade(machine, to_series=to_series)
+        for machine in machines]
+    asyncio.gather(*prepare_group)
+    await prepare_series_upgrade(leader_machine, to_series=to_series)
+    if leader_machine not in completed_machines:
+        machines.append(leader_machine)
+    # do the dist upgrade
+    await dist_upgrades(machines)
+    # do a do-release-upgrade
+    await do_release_upgrades(machines)
+    # do a reboot
+    await reboots(machines)
+
+    await complete_series_upgrade(machines, to_series)
+    if origin:
+        await os_utils.async_set_origin(application, origin)
+
+
+async def serial_series_upgrade(
+        application, from_series='xenial', to_series='bionic',
+        origin='openstack-origin', pause_non_leader_primary=True,
+        pause_non_leader_subordinate=True, post_upgrade_functions=None,
+        completed_machines=None, files=None, workaround_script=None,
+        follower_first=False,):
+    """Perform series upgrade on an application in series.
+
+    :param unit_name: Unit Name
+    :type unit_name: str
+    :param machine_num: Machine number
+    :type machine_num: str
+    :param from_series: The series from which to upgrade
+    :type from_series: str
+    :param to_series: The series to which to upgrade
+    :type to_series: str
+    :param origin: The configuration setting variable name for changing origin
+                   source. (openstack-origin or source)
+    :type origin: str
+    :param files: Workaround files to scp to unit under upgrade
+    :type files: list
+    :param workaround_script: Workaround script to run during series upgrade
+    :type workaround_script: str
+    :param follower_first: Should the follower(s) be upgraded first
+    :type follower_first: bool
+    :returns: None
+    :rtype: None
+    """
+    if completed_machines is None:
+        completed_machines = []
+    if files is None:
+        files = []
+    if post_upgrade_functions is None:
+        post_upgrade_functions = []
+    logging.info(
+        "About to upgrade the units of {} in serial (follower first: {})"
+        .format(application, follower_first))
+    # return
+    status = (await model.async_get_status()).applications[application]
+    leader, non_leaders = await get_leader_and_non_leaders(application)
+    for leader_name, leader_unit in leader.items():
+        leader_machine = leader_unit["machine"]
+        leader = leader_name
+
+    machines = [
+        unit["machine"] for name, unit
+        in non_leaders.items()
+        if unit['machine'] not in completed_machines]
+
+    await maybe_pause_things(
+        status,
+        non_leaders,
+        pause_non_leader_subordinate,
+        pause_non_leader_primary)
+    await series_upgrade_utils.async_set_series(
+        application, to_series=to_series)
+
+    if not follower_first and leader_machine not in completed_machines:
+        await prepare_series_upgrade(leader_machine, to_series)
+        logging.info(
+            "About to dist-upgrade the leader ({})".format(leader_machine))
+        # upgrade the leader
+        await dist_upgrades([leader_machine])
+        # do a do-release-upgrade
+        await do_release_upgrades([leader_machine])
+        # do a reboot
+        await reboots([leader_machine])
+        await complete_series_upgrade([leader_machine], to_series)
+
+    for machine in machines:
+        await prepare_series_upgrade(machine, to_series)
+        logging.debug(
+            "About to dist-upgrade follower {}".format(machine))
+        await dist_upgrades([machine])
+        # do a do-release-upgrade
+        await do_release_upgrades([machine])
+        # do a reboot
+        await reboots([machine])
+        await complete_series_upgrade([machine], to_series)
+
+    if follower_first and leader_machine not in completed_machines:
+        await prepare_series_upgrade(leader_machine, to_series)
+        logging.info(
+            "About to dist-upgrade the leader ({})".format(leader_machine))
+        # upgrade the leader
+        await dist_upgrades([leader_machine])
+        # do a do-release-upgrade
+        await do_release_upgrades([leader_machine])
+        # do a reboot
+        await reboots([leader_machine])
+        await complete_series_upgrade([leader_machine], to_series)
+    if origin:
+        await os_utils.async_set_origin(application, origin)
+
+
+async def maybe_pause_things(
+        status, units, pause_non_leader_subordinate=True,
+        pause_non_leader_primary=True):
+    """Pause the non-leaders, based on the run configuration.
+
+    :param status: Juju status for an application
+    :type status: juju.applications
+    :param units: List of units to paybe pause
+    :type units: LIst[str]
+    :param pause_non_leader_subordinate: Should the non leader
+                                         subordinate be paused
+    :type pause_non_leader_subordinate: bool
+    :param pause_non_leader_primary: Should the non leaders be paused
+    :type pause_non_leader_primary: bool
+    :returns: Nothing
+    :trype: None
+    """
+    subordinate_pauses = []
+    leader_pauses = []
+    for unit in units:
+        if pause_non_leader_subordinate:
+            if status["units"][unit].get("subordinates"):
+                for subordinate in status["units"][unit]["subordinates"]:
+                    _app = subordinate.split('/')[0]
+                    if _app in SUBORDINATE_PAUSE_RESUME_BLACKLIST:
+                        logging.info("Skipping pausing {} - blacklisted"
+                                     .format(subordinate))
+                    else:
+                        logging.info("Pausing {}".format(subordinate))
+                        subordinate_pauses.append(model.async_run_action(
+                            subordinate, "pause", action_params={}))
+        if pause_non_leader_primary:
+            logging.info("Pausing {}".format(unit))
+            leader_pauses.append(
+                model.async_run_action(unit, "pause", action_params={}))
+    await asyncio.gather(*subordinate_pauses)
+    await asyncio.gather(*leader_pauses)
+
+
+async def dist_upgrades(machines):
+    """Run dist-upgrade on unit after update package db.
+
+    :param machines: List of machines to upgrade
+    :type machines: List[str]
+    :returns: None
+    :rtype: None
+    """
+    upgrade_group = []
+    for machine in machines:
+        upgrade_group.append(async_dist_upgrade(machine))
+    logging.info(
+        "About to await the dist upgrades of {}".format(machines))
+    await asyncio.gather(*upgrade_group)
+
+
+async def do_release_upgrades(machines):
+    """Run do-release-upgrade noninteractive.
+
+    :param machines: List of machines to upgrade
+    :type machines: List[str]
+    :returns: None
+    :rtype: None
+    """
+    upgrade_group = []
+    for machine in machines:
+        upgrade_group.append(async_do_release_upgrade(machine))
+    logging.info(
+        "About to await the do release upgrades of {}".format(machines))
+    await asyncio.gather(*upgrade_group)
+
+
+async def reboots(machines):
+    """Reboot all of the listed machines.
+
+    :param machines: A list of machines, ex: ['1', '2']
+    :type machines: List[str]
+    :returns: None
+    :rtype: None
+    """
+    upgrade_group = []
+    for machine in machines:
+        upgrade_group.append(reboot(machine))
+    logging.info("About to await the reboots of {}".format(machines))
+    await asyncio.gather(*upgrade_group)
+
+
+async def get_units(application):
+    """Get all units for the named application.
+
+    :param application: The application to get units of
+    :type application: str
+    :returns: The units for a specified application
+    :rtype: Dict[str, juju.Unit]
+    """
+    status = (await model.async_get_status()).applications[application]
+    return status["units"]
+
+
+async def get_leader_and_non_leaders(application):
+    """Get the leader and non-leader Juju units.
+
+    This function returns a tuple that looks like:
+
+    ({
+        'unit/1': juju.Unit,
+    },
+    {
+        'unit/0': juju.Unit,
+        'unit/2': juju.unit,
+    })
+
+    The first entry of this tuple is the leader, and the second is
+    all non-leader units.
+
+    :param application: Application to fetch details for
+    :type application: str
+    :returns: A tuple of dicts identifying leader and non-leaders
+    :rtype: Dict[str, List[juju.Unit]]
+    """
+    logging.info(
+        "Configuring leader / non leaders for {}".format(application))
+    # if completed_machines is None:
+    #     completed_machines = []
+    leader = None
+    non_leaders = {}
+    for name, unit in (await get_units(application)).items():
+        if unit.get("leader"):
+            leader = {name: unit}
+            # leader = status["units"][unit]["machine"]
+        else:
+            non_leaders[name] = unit
+            # machine = status["units"][unit]["machine"]
+            # if machine not in completed_machines:
+            #     non_leaders.append(machine)
+    return (leader, non_leaders)
+
+
+async def prepare_series_upgrade(machine, to_series):
+    """Execute juju series-upgrade prepare on machine.
+
+    NOTE: This is a new feature in juju behind a feature flag and not yet in
+    libjuju.
+    export JUJU_DEV_FEATURE_FLAGS=upgrade-series
+    :param machine_num: Machine number
+    :type machine_num: str
+    :param to_series: The series to which to upgrade
+    :type to_series: str
+    :returns: None
+    :rtype: None
+    """
+    logging.debug("Preparing series upgrade for: {}".format(machine))
+    await series_upgrade_utils.async_prepare_series_upgrade(
+        machine, to_series=to_series)
+
+
+async def reboot(unit):
+    """Reboot the named machine.
+
+    :param unit: Machine to reboot
+    :type unit: str
+    :returns: Nothing
+    :rtype: None
+    """
+    try:
+        await run_on_machine(unit, 'shutdown --reboot now & exit')
+        # await run_on_machine(unit, "sudo reboot && exit")
+    except subprocess.CalledProcessError as e:
+        logging.warn("Error doing reboot: {}".format(e))
+        pass
+
+
+async def complete_series_upgrade(machines, to_series):
+    """Execute juju series-upgrade complete on machine.
+
+    NOTE: This is a new feature in juju behind a feature flag and not yet in
+    libjuju.
+    export JUJU_DEV_FEATURE_FLAGS=upgrade-series
+    :param machine_num: Machine number
+    :type machine_num: str
+    :returns: None
+    :rtype: None
+    """
+    logging.info("Completing series upgrade for {}".format(machines))
+    group = []
+    for machine in machines:
+        # This can fail on the non-leaders if the leader goes first?
+        group.append(
+            series_upgrade_utils.async_complete_series_upgrade(machine))
+    await asyncio.gather(*group)
+
+
+async def run_on_machine(machine, command, model_name=None, timeout=None):
+    """Juju run on unit.
+
+    :param model_name: Name of model unit is in
+    :type model_name: str
+    :param unit_name: Name of unit to match
+    :type unit: str
+    :param command: Command to execute
+    :type command: str
+    :param timeout: How long in seconds to wait for command to complete
+    :type timeout: int
+    :returns: action.data['results'] {'Code': '', 'Stderr': '', 'Stdout': ''}
+    :rtype: dict
+    """
+    cmd = ['juju', 'run', '--machine={}'.format(machine)]
+    if model_name:
+        cmd.append('--model={}'.format(model_name))
+    if timeout:
+        cmd.append('--timeout={}'.format(timeout))
+    cmd.append(command)
+    logging.debug("About to call '{}'".format(cmd))
+    await os_utils.check_call(cmd)
+
+
+async def async_dist_upgrade(machine):
+    """Run dist-upgrade on unit after update package db.
+
+    :param machine: Machine Number
+    :type machine: str
+    :returns: None
+    :rtype: None
+    """
+    logging.info('Updating package db ' + machine)
+    update_cmd = 'sudo apt update'
+    await run_on_machine(machine, update_cmd)
+
+    logging.info('Updating existing packages ' + machine)
+    dist_upgrade_cmd = (
+        """yes | sudo DEBIAN_FRONTEND=noninteractive apt --assume-yes """
+        """-o "Dpkg::Options::=--force-confdef" """
+        """-o "Dpkg::Options::=--force-confold" dist-upgrade""")
+    await run_on_machine(machine, dist_upgrade_cmd)
+
+
+async def async_do_release_upgrade(machine):
+    """Run do-release-upgrade noninteractive.
+
+    :param machine: Machine Name
+    :type machine: str
+    :returns: None
+    :rtype: None
+    """
+    logging.info('Upgrading ' + machine)
+    do_release_upgrade_cmd = (
+        'yes | sudo DEBIAN_FRONTEND=noninteractive '
+        'do-release-upgrade -f DistUpgradeViewNonInteractive')
+
+    await run_on_machine(machine, do_release_upgrade_cmd, timeout="120m")
diff --git a/zaza/openstack/utilities/upgrade_utils.py b/zaza/openstack/utilities/upgrade_utils.py
index 062e9b1..b24aec7 100644
--- a/zaza/openstack/utilities/upgrade_utils.py
+++ b/zaza/openstack/utilities/upgrade_utils.py
@@ -20,15 +20,17 @@ import zaza.model
 
 
 SERVICE_GROUPS = collections.OrderedDict([
+    ('Stateful Services', ['percona-cluster', 'rabbitmq-server']),
     ('Core Identity', ['keystone']),
     ('Control Plane', [
-        'aodh', 'barbican', 'ceilometer', 'ceph-mon', 'ceph-fs',
+        'aodh', 'barbican', 'ceilometer', 'ceph-fs',
         'ceph-radosgw', 'cinder', 'designate',
         'designate-bind', 'glance', 'gnocchi', 'heat', 'manila',
         'manila-generic', 'neutron-api', 'neutron-gateway', 'placement',
         'nova-cloud-controller', 'openstack-dashboard']),
     ('Data Plane', [
-        'nova-compute', 'ceph-osd', 'swift-proxy', 'swift-storage'])
+        'nova-compute', 'ceph-mon', 'ceph-osd',
+        'swift-proxy', 'swift-storage'])
 ])
 
 UPGRADE_EXCLUDE_LIST = ['rabbitmq-server', 'percona-cluster']

From 186708e88353bbb99819c6ce695a4315791337cf Mon Sep 17 00:00:00 2001
From: Chris MacNaughton <chris.macnaughton@canonical.com>
Date: Thu, 9 Apr 2020 09:12:14 +0200
Subject: [PATCH 02/19] Change upgrade process to handle each machine fully

Updating each machine in a fully self-contained way makes
it easy to add per-machine and (possibly) per-application
post-upgrade function calls. This still performs the actual
series upgrades fully in parallel but will, for example,
start the do-release-upgrade on one unit while another is
still performing the initial dist-upgrade
---
 .../series_upgrade/parallel_tests.py          |   2 +-
 .../utilities/parallel_series_upgrade.py      | 113 ++++++++++++------
 2 files changed, 80 insertions(+), 35 deletions(-)

diff --git a/zaza/openstack/charm_tests/series_upgrade/parallel_tests.py b/zaza/openstack/charm_tests/series_upgrade/parallel_tests.py
index 1b64994..7c9dba0 100644
--- a/zaza/openstack/charm_tests/series_upgrade/parallel_tests.py
+++ b/zaza/openstack/charm_tests/series_upgrade/parallel_tests.py
@@ -28,7 +28,7 @@ from zaza.openstack.utilities import (
     upgrade_utils as upgrade_utils,
 )
 from zaza.openstack.charm_tests.nova.tests import LTSGuestCreateTest
-from zaza.openstack.utilities.parallel_series_upgrade import (
+from zaza.openstack.utilities import (
     parallel_series_upgrade,
 )
 
diff --git a/zaza/openstack/utilities/parallel_series_upgrade.py b/zaza/openstack/utilities/parallel_series_upgrade.py
index a959838..1585ef7 100755
--- a/zaza/openstack/utilities/parallel_series_upgrade.py
+++ b/zaza/openstack/utilities/parallel_series_upgrade.py
@@ -125,6 +125,16 @@ async def parallel_series_upgrade(
     :param origin: The configuration setting variable name for changing origin
                    source. (openstack-origin or source)
     :type origin: str
+    :param pause_non_leader_primary: Whether the non-leader applications should
+                                     be paused
+    :type pause_non_leader_primary: bool
+    :param pause_non_leader_subordinate: Whether the non-leader subordinate
+                                         hacluster applications should be
+                                         paused
+    :type pause_non_leader_subordinate: bool
+    :param post_upgrade_functions: A list of Zaza functions to call when
+                                   the upgrade is complete on each machine
+    :type post_upgrade_functions: List[str]
     :param files: Workaround files to scp to unit under upgrade
     :type files: list
     :param workaround_script: Workaround script to run during series upgrade
@@ -170,14 +180,14 @@ async def parallel_series_upgrade(
     await prepare_series_upgrade(leader_machine, to_series=to_series)
     if leader_machine not in completed_machines:
         machines.append(leader_machine)
-    # do the dist upgrade
-    await dist_upgrades(machines)
-    # do a do-release-upgrade
-    await do_release_upgrades(machines)
-    # do a reboot
-    await reboots(machines)
-
-    await complete_series_upgrade(machines, to_series)
+    upgrade_group = [
+        series_upgrade_machine(
+            machine, from_series=from_series, to_series=to_series,
+            files=files, workaround_script=None,
+            post_upgrade_functions=post_upgrade_functions)
+        for machine in machines
+    ]
+    await asyncio.gather(*upgrade_group)
     if origin:
         await os_utils.async_set_origin(application, origin)
 
@@ -201,6 +211,16 @@ async def serial_series_upgrade(
     :param origin: The configuration setting variable name for changing origin
                    source. (openstack-origin or source)
     :type origin: str
+    :param pause_non_leader_primary: Whether the non-leader applications should
+                                     be paused
+    :type pause_non_leader_primary: bool
+    :param pause_non_leader_subordinate: Whether the non-leader subordinate
+                                         hacluster applications should be
+                                         paused
+    :type pause_non_leader_subordinate: bool
+    :param post_upgrade_functions: A list of Zaza functions to call when
+                                   the upgrade is complete on each machine
+    :type post_upgrade_functions: List[str]
     :param files: Workaround files to scp to unit under upgrade
     :type files: list
     :param workaround_script: Workaround script to run during series upgrade
@@ -241,42 +261,67 @@ async def serial_series_upgrade(
 
     if not follower_first and leader_machine not in completed_machines:
         await prepare_series_upgrade(leader_machine, to_series)
-        logging.info(
-            "About to dist-upgrade the leader ({})".format(leader_machine))
-        # upgrade the leader
-        await dist_upgrades([leader_machine])
-        # do a do-release-upgrade
-        await do_release_upgrades([leader_machine])
-        # do a reboot
-        await reboots([leader_machine])
-        await complete_series_upgrade([leader_machine], to_series)
+        logging.info("About to upgrade leader of {}: {}"
+                     .format(application, leader_machine))
+        await series_upgrade_machine(
+            leader_machine, from_series=from_series, to_series=to_series,
+            files=files, workaround_script=None,
+            post_upgrade_functions=post_upgrade_functions)
 
     for machine in machines:
         await prepare_series_upgrade(machine, to_series)
-        logging.debug(
-            "About to dist-upgrade follower {}".format(machine))
-        await dist_upgrades([machine])
-        # do a do-release-upgrade
-        await do_release_upgrades([machine])
-        # do a reboot
-        await reboots([machine])
-        await complete_series_upgrade([machine], to_series)
+        logging.info("About to upgrade follower of {}: {}"
+                     .format(application, machine))
+        await series_upgrade_machine(
+            machine, from_series=from_series, to_series=to_series,
+            files=files, workaround_script=None,
+            post_upgrade_functions=post_upgrade_functions)
 
     if follower_first and leader_machine not in completed_machines:
         await prepare_series_upgrade(leader_machine, to_series)
-        logging.info(
-            "About to dist-upgrade the leader ({})".format(leader_machine))
-        # upgrade the leader
-        await dist_upgrades([leader_machine])
-        # do a do-release-upgrade
-        await do_release_upgrades([leader_machine])
-        # do a reboot
-        await reboots([leader_machine])
-        await complete_series_upgrade([leader_machine], to_series)
+        logging.info("About to upgrade leader of {}: {}"
+                     .format(application, leader_machine))
+        await series_upgrade_machine(
+            leader_machine, from_series=from_series, to_series=to_series,
+            files=files, workaround_script=None,
+            post_upgrade_functions=post_upgrade_functions)
     if origin:
         await os_utils.async_set_origin(application, origin)
 
 
+async def series_upgrade_machine(
+        machine, from_series='xenial', to_series='bionic',
+        files=None, workaround_script=None, post_upgrade_functions=None):
+    """Perform series upgrade on an machine.
+
+    :param machine_num: Machine number
+    :type machine_num: str
+    :param from_series: The series from which to upgrade
+    :type from_series: str
+    :param to_series: The series to which to upgrade
+    :type to_series: str
+    :param files: Workaround files to scp to unit under upgrade
+    :type files: list
+    :param workaround_script: Workaround script to run during series upgrade
+    :type workaround_script: str
+    :param post_upgrade_functions: A list of Zaza functions to call when
+                                   the upgrade is complete on each machine
+    :type post_upgrade_functions: List[str]
+    :returns: None
+    :rtype: None
+    """
+    logging.info(
+        "About to dist-upgrade ({})".format(machine))
+    # upgrade the do the dist upgrade
+    await async_dist_upgrade(machine)
+    # do a do-release-upgrade
+    await async_do_release_upgrade(machine)
+    # do a reboot
+    await reboot(machine)
+    await series_upgrade_utils.async_complete_series_upgrade(machine)
+    series_upgrade_utils.run_post_upgrade_functions(post_upgrade_functions)
+
+
 async def maybe_pause_things(
         status, units, pause_non_leader_subordinate=True,
         pause_non_leader_primary=True):

From 2a98274d67589a5a251c0fa9082f28e27abc6833 Mon Sep 17 00:00:00 2001
From: Chris MacNaughton <chris.macnaughton@canonical.com>
Date: Thu, 9 Apr 2020 10:56:31 +0200
Subject: [PATCH 03/19] Add a few more function hooks to the series upgrade
 process

By including pre and pos-application functions, charm
series upgrades can be handled in a more generic way,
even when they require running additional actions
before a unit is upgraded, or after the whole application
is upgraded.
---
 zaza/openstack/charm_tests/mysql/utils.py     | 26 +++++
 .../charm_tests/rabbitmq_server/utils.py      |  8 ++
 .../series_upgrade/parallel_tests.py          | 16 ----
 .../utilities/parallel_series_upgrade.py      | 94 +++++++++++++++++--
 zaza/openstack/utilities/upgrade_utils.py     |  4 +-
 5 files changed, 123 insertions(+), 25 deletions(-)
 create mode 100644 zaza/openstack/charm_tests/mysql/utils.py

diff --git a/zaza/openstack/charm_tests/mysql/utils.py b/zaza/openstack/charm_tests/mysql/utils.py
new file mode 100644
index 0000000..2622903
--- /dev/null
+++ b/zaza/openstack/charm_tests/mysql/utils.py
@@ -0,0 +1,26 @@
+# Copyright 2020 Canonical Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Module of functions for interfacing with the percona-cluster charm."""
+
+import zaza.model as model
+
+
+def complete_cluster_series_upgrade():
+    """Run the complete-cluster-series-upgrade action on the lead unit."""
+    # TODO: Make this work across either mysql or percona-cluster names
+    model.run_action_on_leader(
+        'mysql',
+        'complete-cluster-series-upgrade',
+        action_params={})
diff --git a/zaza/openstack/charm_tests/rabbitmq_server/utils.py b/zaza/openstack/charm_tests/rabbitmq_server/utils.py
index 9826add..62db7e6 100644
--- a/zaza/openstack/charm_tests/rabbitmq_server/utils.py
+++ b/zaza/openstack/charm_tests/rabbitmq_server/utils.py
@@ -584,3 +584,11 @@ def _post_check_unit_cluster_nodes(unit, nodes, unit_node_names):
                           'after unit removal.\n'
                           ''.format(unit_name, node))
     return errors
+
+
+def complete_cluster_series_upgrade():
+    """Run the complete-cluster-series-upgrade action on the lead unit."""
+    zaza.model.run_action_on_leader(
+        'rabbitmq-server',
+        'complete-cluster-series-upgrade',
+        action_params={})
diff --git a/zaza/openstack/charm_tests/series_upgrade/parallel_tests.py b/zaza/openstack/charm_tests/series_upgrade/parallel_tests.py
index 7c9dba0..5307339 100644
--- a/zaza/openstack/charm_tests/series_upgrade/parallel_tests.py
+++ b/zaza/openstack/charm_tests/series_upgrade/parallel_tests.py
@@ -101,22 +101,6 @@ class ParallelSeriesUpgradeTest(unittest.TestCase):
                         files=files))
             asyncio.get_event_loop().run_until_complete(
                 asyncio.gather(*upgrade_functions))
-
-            if "rabbitmq-server" in apps:
-                logging.info(
-                    "Running complete-cluster-series-upgrade action on leader")
-                model.run_action_on_leader(
-                    'rabbitmq-server',
-                    'complete-cluster-series-upgrade',
-                    action_params={})
-
-            if "percona-cluster" in apps:
-                logging.info(
-                    "Running complete-cluster-series-upgrade action on leader")
-                model.run_action_on_leader(
-                    'mysql',
-                    'complete-cluster-series-upgrade',
-                    action_params={})
             model.block_until_all_units_idle()
             logging.info("Finished {}".format(group_name))
         logging.info("Done!")
diff --git a/zaza/openstack/utilities/parallel_series_upgrade.py b/zaza/openstack/utilities/parallel_series_upgrade.py
index 1585ef7..622c5ac 100755
--- a/zaza/openstack/utilities/parallel_series_upgrade.py
+++ b/zaza/openstack/utilities/parallel_series_upgrade.py
@@ -23,6 +23,7 @@ import logging
 import subprocess
 
 from zaza import model
+from zaza.charm_lifecycle import utils as cl_utils
 import zaza.openstack.utilities.generic as os_utils
 import zaza.openstack.utilities.series_upgrade as series_upgrade_utils
 from zaza.openstack.utilities.series_upgrade import (
@@ -45,6 +46,8 @@ def app_config(charm_name):
         'pause_non_leader_subordinate': True,
         'pause_non_leader_primary': True,
         'post_upgrade_functions': [],
+        'pre_upgrade_functions': [],
+        'post_application_upgrade_functions': [],
         'follower_first': False, }
     _app_settings = collections.defaultdict(lambda: default)
     ceph = {
@@ -55,11 +58,25 @@ def app_config(charm_name):
     exceptions = {
         'rabbitmq-server': {
             'origin': 'source',
-            'pause_non_leader_subordinate': False, },
-        'percona-cluster': {'origin': 'source', },
+            'pause_non_leader_subordinate': False,
+            'post_application_upgrade_functions': [
+                ('zaza.openstack.charm_tests.rabbitmq_server.utils.'
+                 'complete_cluster_series_upgrade')]
+        },
+        'percona-cluster': {
+            'origin': 'source',
+            'post_application_upgrade_functions': [
+                ('zaza.openstack.charm_tests.mysql.utils.'
+                 'complete_cluster_series_upgrade')]
+        },
         'nova-compute': {
             'pause_non_leader_primary': False,
-            'pause_non_leader_subordinate': False, },
+            'pause_non_leader_subordinate': False,
+            # TODO
+            # 'pre_upgrade_functions': [
+            #     'zaza.openstack.charm_tests.nova_compute.setup.evacuate'
+            # ]
+        },
         'ceph': ceph,
         'ceph-mon': ceph,
         'ceph-osd': ceph,
@@ -111,7 +128,8 @@ async def parallel_series_upgrade(
         origin='openstack-origin', pause_non_leader_primary=True,
         pause_non_leader_subordinate=True, post_upgrade_functions=None,
         completed_machines=None, files=None, workaround_script=None,
-        follower_first=False):
+        follower_first=False, pre_upgrade_functions=None,
+        post_application_upgrade_functions=None):
     """Perform series upgrade on an application in parallel.
 
     :param unit_name: Unit Name
@@ -132,9 +150,17 @@ async def parallel_series_upgrade(
                                          hacluster applications should be
                                          paused
     :type pause_non_leader_subordinate: bool
+
+    :param pre_upgrade_functions: A list of Zaza functions to call before
+                                  the upgrade is started on each machine
+    :type pre_upgrade_functions: List[str]
     :param post_upgrade_functions: A list of Zaza functions to call when
                                    the upgrade is complete on each machine
     :type post_upgrade_functions: List[str]
+    :param post_application_upgrade_functions: A list of Zaza functions
+                                   to call when the upgrade is complete
+                                   on all machine in the application
+    :type post_application_upgrade_functions: List[str]
     :param files: Workaround files to scp to unit under upgrade
     :type files: list
     :param workaround_script: Workaround script to run during series upgrade
@@ -150,12 +176,16 @@ async def parallel_series_upgrade(
         files = []
     if post_upgrade_functions is None:
         post_upgrade_functions = []
+    if pre_upgrade_functions is None:
+        pre_upgrade_functions = []
+    if post_application_upgrade_functions is None:
+        post_application_upgrade_functions = []
     if follower_first:
         logging.error("leader_first is ignored for parallel upgrade")
     logging.info(
         "About to upgrade the units of {} in parallel (follower first: {})"
         .format(application, follower_first))
-    # return
+
     status = (await model.async_get_status()).applications[application]
     leader, non_leaders = await get_leader_and_non_leaders(application)
     for leader_name, leader_unit in leader.items():
@@ -165,6 +195,7 @@ async def parallel_series_upgrade(
         unit["machine"] for name, unit
         in non_leaders.items()
         if unit['machine'] not in completed_machines]
+
     await maybe_pause_things(
         status,
         non_leaders,
@@ -190,6 +221,7 @@ async def parallel_series_upgrade(
     await asyncio.gather(*upgrade_group)
     if origin:
         await os_utils.async_set_origin(application, origin)
+    run_post_application_upgrade_functions(post_application_upgrade_functions)
 
 
 async def serial_series_upgrade(
@@ -197,7 +229,8 @@ async def serial_series_upgrade(
         origin='openstack-origin', pause_non_leader_primary=True,
         pause_non_leader_subordinate=True, post_upgrade_functions=None,
         completed_machines=None, files=None, workaround_script=None,
-        follower_first=False,):
+        follower_first=False, pre_upgrade_functions=None,
+        post_application_upgrade_functions=None):
     """Perform series upgrade on an application in series.
 
     :param unit_name: Unit Name
@@ -218,9 +251,16 @@ async def serial_series_upgrade(
                                          hacluster applications should be
                                          paused
     :type pause_non_leader_subordinate: bool
+    :param pre_upgrade_functions: A list of Zaza functions to call before
+                                  the upgrade is started on each machine
+    :type pre_upgrade_functions: List[str]
     :param post_upgrade_functions: A list of Zaza functions to call when
                                    the upgrade is complete on each machine
     :type post_upgrade_functions: List[str]
+    :param post_application_upgrade_functions: A list of Zaza functions
+                                   to call when the upgrade is complete
+                                   on all machine in the application
+    :type post_application_upgrade_functions: List[str]
     :param files: Workaround files to scp to unit under upgrade
     :type files: list
     :param workaround_script: Workaround script to run during series upgrade
@@ -236,6 +276,10 @@ async def serial_series_upgrade(
         files = []
     if post_upgrade_functions is None:
         post_upgrade_functions = []
+    if pre_upgrade_functions is None:
+        pre_upgrade_functions = []
+    if post_application_upgrade_functions is None:
+        post_application_upgrade_functions = []
     logging.info(
         "About to upgrade the units of {} in serial (follower first: {})"
         .format(application, follower_first))
@@ -287,11 +331,13 @@ async def serial_series_upgrade(
             post_upgrade_functions=post_upgrade_functions)
     if origin:
         await os_utils.async_set_origin(application, origin)
+    run_post_application_upgrade_functions(post_application_upgrade_functions)
 
 
 async def series_upgrade_machine(
         machine, from_series='xenial', to_series='bionic',
-        files=None, workaround_script=None, post_upgrade_functions=None):
+        files=None, workaround_script=None, post_upgrade_functions=None,
+        pre_upgrade_functions=None):
     """Perform series upgrade on an machine.
 
     :param machine_num: Machine number
@@ -304,6 +350,9 @@ async def series_upgrade_machine(
     :type files: list
     :param workaround_script: Workaround script to run during series upgrade
     :type workaround_script: str
+    :param pre_upgrade_functions: A list of Zaza functions to call before
+                                  the upgrade is started on each machine
+    :type pre_upgrade_functions: List[str]
     :param post_upgrade_functions: A list of Zaza functions to call when
                                    the upgrade is complete on each machine
     :type post_upgrade_functions: List[str]
@@ -312,6 +361,8 @@ async def series_upgrade_machine(
     """
     logging.info(
         "About to dist-upgrade ({})".format(machine))
+
+    run_pre_upgrade_functions(pre_upgrade_functions)
     # upgrade the do the dist upgrade
     await async_dist_upgrade(machine)
     # do a do-release-upgrade
@@ -322,6 +373,35 @@ async def series_upgrade_machine(
     series_upgrade_utils.run_post_upgrade_functions(post_upgrade_functions)
 
 
+def run_pre_upgrade_functions(machine, pre_upgrade_functions):
+    """Execute list supplied functions.
+
+    Each of the supplied functions will be called with a single
+    argument of the machine that is about to be upgraded.
+
+    :param machine: Machine that is about to be upgraded
+    :type machine: str
+    :param pre_upgrade_functions: List of functions
+    :type pre_upgrade_functions: [function, function, ...]
+    """
+    if pre_upgrade_functions:
+        for func in pre_upgrade_functions:
+            logging.info("Running {}".format(func))
+            cl_utils.get_class(func)(machine)
+
+
+def run_post_application_upgrade_functions(post_upgrade_functions):
+    """Execute list supplied functions.
+
+    :param post_upgrade_functions: List of functions
+    :type post_upgrade_functions: [function, function, ...]
+    """
+    if post_upgrade_functions:
+        for func in post_upgrade_functions:
+            logging.info("Running {}".format(func))
+            cl_utils.get_class(func)()
+
+
 async def maybe_pause_things(
         status, units, pause_non_leader_subordinate=True,
         pause_non_leader_primary=True):
diff --git a/zaza/openstack/utilities/upgrade_utils.py b/zaza/openstack/utilities/upgrade_utils.py
index b24aec7..c864b9b 100644
--- a/zaza/openstack/utilities/upgrade_utils.py
+++ b/zaza/openstack/utilities/upgrade_utils.py
@@ -21,7 +21,7 @@ import zaza.model
 
 SERVICE_GROUPS = collections.OrderedDict([
     ('Stateful Services', ['percona-cluster', 'rabbitmq-server']),
-    ('Core Identity', ['keystone']),
+    ('Core Identity', ['keystone', 'ceph-mon']),
     ('Control Plane', [
         'aodh', 'barbican', 'ceilometer', 'ceph-fs',
         'ceph-radosgw', 'cinder', 'designate',
@@ -29,7 +29,7 @@ SERVICE_GROUPS = collections.OrderedDict([
         'manila-generic', 'neutron-api', 'neutron-gateway', 'placement',
         'nova-cloud-controller', 'openstack-dashboard']),
     ('Data Plane', [
-        'nova-compute', 'ceph-mon', 'ceph-osd',
+        'nova-compute', 'ceph-osd',
         'swift-proxy', 'swift-storage'])
 ])
 

From 392470e946d45caa3a3f503bd805727f883ff112 Mon Sep 17 00:00:00 2001
From: Chris MacNaughton <chris.macnaughton@canonical.com>
Date: Thu, 9 Apr 2020 15:05:53 +0200
Subject: [PATCH 04/19] Cleanup based on feedback

This change removes some redundant checking at setup
time and splits the long parameter lists into one
argument per line lists
---
 .../utilities/parallel_series_upgrade.py      | 142 ++++++------------
 1 file changed, 47 insertions(+), 95 deletions(-)

diff --git a/zaza/openstack/utilities/parallel_series_upgrade.py b/zaza/openstack/utilities/parallel_series_upgrade.py
index 622c5ac..e08ede9 100755
--- a/zaza/openstack/utilities/parallel_series_upgrade.py
+++ b/zaza/openstack/utilities/parallel_series_upgrade.py
@@ -124,12 +124,20 @@ def upgrade_ubuntu_lite(from_series='xenial', to_series='bionic'):
 
 
 async def parallel_series_upgrade(
-        application, from_series='xenial', to_series='bionic',
-        origin='openstack-origin', pause_non_leader_primary=True,
-        pause_non_leader_subordinate=True, post_upgrade_functions=None,
-        completed_machines=None, files=None, workaround_script=None,
-        follower_first=False, pre_upgrade_functions=None,
-        post_application_upgrade_functions=None):
+    application,
+    from_series='xenial',
+    to_series='bionic',
+    origin='openstack-origin',
+    pause_non_leader_primary=True,
+    pause_non_leader_subordinate=True,
+    pre_upgrade_functions=None,
+    post_upgrade_functions=None,
+    post_application_upgrade_functions=None,
+    completed_machines=None,
+    follower_first=False,
+    files=None,
+    workaround_script=None
+):
     """Perform series upgrade on an application in parallel.
 
     :param unit_name: Unit Name
@@ -160,26 +168,18 @@ async def parallel_series_upgrade(
     :param post_application_upgrade_functions: A list of Zaza functions
                                    to call when the upgrade is complete
                                    on all machine in the application
+    :param follower_first: Should the follower(s) be upgraded first
+    :type follower_first: bool
     :type post_application_upgrade_functions: List[str]
     :param files: Workaround files to scp to unit under upgrade
     :type files: list
     :param workaround_script: Workaround script to run during series upgrade
     :type workaround_script: str
-    :param follower_first: Should the follower(s) be upgraded first
-    :type follower_first: bool
     :returns: None
     :rtype: None
     """
     if completed_machines is None:
         completed_machines = []
-    if files is None:
-        files = []
-    if post_upgrade_functions is None:
-        post_upgrade_functions = []
-    if pre_upgrade_functions is None:
-        pre_upgrade_functions = []
-    if post_application_upgrade_functions is None:
-        post_application_upgrade_functions = []
     if follower_first:
         logging.error("leader_first is ignored for parallel upgrade")
     logging.info(
@@ -213,8 +213,8 @@ async def parallel_series_upgrade(
         machines.append(leader_machine)
     upgrade_group = [
         series_upgrade_machine(
-            machine, from_series=from_series, to_series=to_series,
-            files=files, workaround_script=None,
+            machine,
+            files=files, workaround_script=workaround_script,
             post_upgrade_functions=post_upgrade_functions)
         for machine in machines
     ]
@@ -225,13 +225,21 @@ async def parallel_series_upgrade(
 
 
 async def serial_series_upgrade(
-        application, from_series='xenial', to_series='bionic',
-        origin='openstack-origin', pause_non_leader_primary=True,
-        pause_non_leader_subordinate=True, post_upgrade_functions=None,
-        completed_machines=None, files=None, workaround_script=None,
-        follower_first=False, pre_upgrade_functions=None,
-        post_application_upgrade_functions=None):
-    """Perform series upgrade on an application in series.
+    application,
+    from_series='xenial',
+    to_series='bionic',
+    origin='openstack-origin',
+    pause_non_leader_primary=True,
+    pause_non_leader_subordinate=True,
+    pre_upgrade_functions=None,
+    post_upgrade_functions=None,
+    post_application_upgrade_functions=None,
+    completed_machines=None,
+    follower_first=False,
+    files=None,
+    workaround_script=None
+):
+    """Perform series upgrade on an application in parallel.
 
     :param unit_name: Unit Name
     :type unit_name: str
@@ -251,6 +259,7 @@ async def serial_series_upgrade(
                                          hacluster applications should be
                                          paused
     :type pause_non_leader_subordinate: bool
+
     :param pre_upgrade_functions: A list of Zaza functions to call before
                                   the upgrade is started on each machine
     :type pre_upgrade_functions: List[str]
@@ -260,26 +269,18 @@ async def serial_series_upgrade(
     :param post_application_upgrade_functions: A list of Zaza functions
                                    to call when the upgrade is complete
                                    on all machine in the application
+    :param follower_first: Should the follower(s) be upgraded first
+    :type follower_first: bool
     :type post_application_upgrade_functions: List[str]
     :param files: Workaround files to scp to unit under upgrade
     :type files: list
     :param workaround_script: Workaround script to run during series upgrade
     :type workaround_script: str
-    :param follower_first: Should the follower(s) be upgraded first
-    :type follower_first: bool
     :returns: None
     :rtype: None
     """
     if completed_machines is None:
         completed_machines = []
-    if files is None:
-        files = []
-    if post_upgrade_functions is None:
-        post_upgrade_functions = []
-    if pre_upgrade_functions is None:
-        pre_upgrade_functions = []
-    if post_application_upgrade_functions is None:
-        post_application_upgrade_functions = []
     logging.info(
         "About to upgrade the units of {} in serial (follower first: {})"
         .format(application, follower_first))
@@ -308,8 +309,8 @@ async def serial_series_upgrade(
         logging.info("About to upgrade leader of {}: {}"
                      .format(application, leader_machine))
         await series_upgrade_machine(
-            leader_machine, from_series=from_series, to_series=to_series,
-            files=files, workaround_script=None,
+            leader_machine,
+            files=files, workaround_script=workaround_script,
             post_upgrade_functions=post_upgrade_functions)
 
     for machine in machines:
@@ -317,8 +318,8 @@ async def serial_series_upgrade(
         logging.info("About to upgrade follower of {}: {}"
                      .format(application, machine))
         await series_upgrade_machine(
-            machine, from_series=from_series, to_series=to_series,
-            files=files, workaround_script=None,
+            machine,
+            files=files, workaround_script=workaround_script,
             post_upgrade_functions=post_upgrade_functions)
 
     if follower_first and leader_machine not in completed_machines:
@@ -326,8 +327,8 @@ async def serial_series_upgrade(
         logging.info("About to upgrade leader of {}: {}"
                      .format(application, leader_machine))
         await series_upgrade_machine(
-            leader_machine, from_series=from_series, to_series=to_series,
-            files=files, workaround_script=None,
+            leader_machine,
+            files=files, workaround_script=workaround_script,
             post_upgrade_functions=post_upgrade_functions)
     if origin:
         await os_utils.async_set_origin(application, origin)
@@ -335,17 +336,15 @@ async def serial_series_upgrade(
 
 
 async def series_upgrade_machine(
-        machine, from_series='xenial', to_series='bionic',
-        files=None, workaround_script=None, post_upgrade_functions=None,
-        pre_upgrade_functions=None):
+        machine,
+        post_upgrade_functions=None,
+        pre_upgrade_functions=None,
+        files=None,
+        workaround_script=None):
     """Perform series upgrade on an machine.
 
     :param machine_num: Machine number
     :type machine_num: str
-    :param from_series: The series from which to upgrade
-    :type from_series: str
-    :param to_series: The series to which to upgrade
-    :type to_series: str
     :param files: Workaround files to scp to unit under upgrade
     :type files: list
     :param workaround_script: Workaround script to run during series upgrade
@@ -441,53 +440,6 @@ async def maybe_pause_things(
     await asyncio.gather(*leader_pauses)
 
 
-async def dist_upgrades(machines):
-    """Run dist-upgrade on unit after update package db.
-
-    :param machines: List of machines to upgrade
-    :type machines: List[str]
-    :returns: None
-    :rtype: None
-    """
-    upgrade_group = []
-    for machine in machines:
-        upgrade_group.append(async_dist_upgrade(machine))
-    logging.info(
-        "About to await the dist upgrades of {}".format(machines))
-    await asyncio.gather(*upgrade_group)
-
-
-async def do_release_upgrades(machines):
-    """Run do-release-upgrade noninteractive.
-
-    :param machines: List of machines to upgrade
-    :type machines: List[str]
-    :returns: None
-    :rtype: None
-    """
-    upgrade_group = []
-    for machine in machines:
-        upgrade_group.append(async_do_release_upgrade(machine))
-    logging.info(
-        "About to await the do release upgrades of {}".format(machines))
-    await asyncio.gather(*upgrade_group)
-
-
-async def reboots(machines):
-    """Reboot all of the listed machines.
-
-    :param machines: A list of machines, ex: ['1', '2']
-    :type machines: List[str]
-    :returns: None
-    :rtype: None
-    """
-    upgrade_group = []
-    for machine in machines:
-        upgrade_group.append(reboot(machine))
-    logging.info("About to await the reboots of {}".format(machines))
-    await asyncio.gather(*upgrade_group)
-
-
 async def get_units(application):
     """Get all units for the named application.
 

From 9cd4b32aa3de43830107239976a7722756c3b8f8 Mon Sep 17 00:00:00 2001
From: Chris MacNaughton <chris.macnaughton@canonical.com>
Date: Thu, 9 Apr 2020 15:08:36 +0200
Subject: [PATCH 05/19] Remove leftover development comments, useless function

---
 .../utilities/parallel_series_upgrade.py      | 28 ++-----------------
 1 file changed, 3 insertions(+), 25 deletions(-)

diff --git a/zaza/openstack/utilities/parallel_series_upgrade.py b/zaza/openstack/utilities/parallel_series_upgrade.py
index e08ede9..87deb51 100755
--- a/zaza/openstack/utilities/parallel_series_upgrade.py
+++ b/zaza/openstack/utilities/parallel_series_upgrade.py
@@ -187,7 +187,7 @@ async def parallel_series_upgrade(
         .format(application, follower_first))
 
     status = (await model.async_get_status()).applications[application]
-    leader, non_leaders = await get_leader_and_non_leaders(application)
+    leader, non_leaders = await get_leader_and_non_leaders(status, application)
     for leader_name, leader_unit in leader.items():
         leader_machine = leader_unit["machine"]
         leader = leader_name
@@ -284,7 +284,6 @@ async def serial_series_upgrade(
     logging.info(
         "About to upgrade the units of {} in serial (follower first: {})"
         .format(application, follower_first))
-    # return
     status = (await model.async_get_status()).applications[application]
     leader, non_leaders = await get_leader_and_non_leaders(application)
     for leader_name, leader_unit in leader.items():
@@ -362,11 +361,8 @@ async def series_upgrade_machine(
         "About to dist-upgrade ({})".format(machine))
 
     run_pre_upgrade_functions(pre_upgrade_functions)
-    # upgrade the do the dist upgrade
     await async_dist_upgrade(machine)
-    # do a do-release-upgrade
     await async_do_release_upgrade(machine)
-    # do a reboot
     await reboot(machine)
     await series_upgrade_utils.async_complete_series_upgrade(machine)
     series_upgrade_utils.run_post_upgrade_functions(post_upgrade_functions)
@@ -440,19 +436,7 @@ async def maybe_pause_things(
     await asyncio.gather(*leader_pauses)
 
 
-async def get_units(application):
-    """Get all units for the named application.
-
-    :param application: The application to get units of
-    :type application: str
-    :returns: The units for a specified application
-    :rtype: Dict[str, juju.Unit]
-    """
-    status = (await model.async_get_status()).applications[application]
-    return status["units"]
-
-
-async def get_leader_and_non_leaders(application):
+async def get_leader_and_non_leaders(status, application):
     """Get the leader and non-leader Juju units.
 
     This function returns a tuple that looks like:
@@ -475,19 +459,13 @@ async def get_leader_and_non_leaders(application):
     """
     logging.info(
         "Configuring leader / non leaders for {}".format(application))
-    # if completed_machines is None:
-    #     completed_machines = []
     leader = None
     non_leaders = {}
-    for name, unit in (await get_units(application)).items():
+    for name, unit in status["units"].items():
         if unit.get("leader"):
             leader = {name: unit}
-            # leader = status["units"][unit]["machine"]
         else:
             non_leaders[name] = unit
-            # machine = status["units"][unit]["machine"]
-            # if machine not in completed_machines:
-            #     non_leaders.append(machine)
     return (leader, non_leaders)
 
 

From 74652f2523bd77d9411e7f08f0a97abe97faabd8 Mon Sep 17 00:00:00 2001
From: Chris MacNaughton <chris.macnaughton@canonical.com>
Date: Thu, 9 Apr 2020 17:29:05 +0200
Subject: [PATCH 06/19] Add unit test coverage for series upgrades

This also includes some tidying
---
 ..._zaza_utilities_parallel_series_upgrade.py | 382 ++++++++++++++++++
 .../utilities/parallel_series_upgrade.py      |  59 +--
 zaza/openstack/utilities/upgrade_utils.py     |   4 +-
 3 files changed, 406 insertions(+), 39 deletions(-)
 create mode 100644 unit_tests/utilities/test_zaza_utilities_parallel_series_upgrade.py

diff --git a/unit_tests/utilities/test_zaza_utilities_parallel_series_upgrade.py b/unit_tests/utilities/test_zaza_utilities_parallel_series_upgrade.py
new file mode 100644
index 0000000..966b54e
--- /dev/null
+++ b/unit_tests/utilities/test_zaza_utilities_parallel_series_upgrade.py
@@ -0,0 +1,382 @@
+# Copyright 2020 Canonical Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import mock
+import unit_tests.utils as ut_utils
+import zaza.openstack.utilities.generic as generic_utils
+import zaza.openstack.utilities.series_upgrade as series_upgrade
+import zaza.openstack.utilities.parallel_series_upgrade as upgrade_utils
+
+FAKE_STATUS = {
+    'can-upgrade-to': '',
+    'charm': 'local:trusty/app-136',
+    'subordinate-to': [],
+    'units': {'app/0': {'leader': True,
+                        'machine': '0',
+                        'subordinates': {
+                            'app-hacluster/0': {
+                                   'charm': 'local:trusty/hacluster-0',
+                                   'leader': True}}},
+              'app/1': {'machine': '1',
+                        'subordinates': {
+                            'app-hacluster/1': {
+                                   'charm': 'local:trusty/hacluster-0'}}},
+              'app/2': {'machine': '2',
+                        'subordinates': {
+                            'app-hacluster/2': {
+                                   'charm': 'local:trusty/hacluster-0'}}}}}
+
+
+class Test_ParallelSeriesUpgradeSync(ut_utils.BaseTestCase):
+    def setUp(self):
+        super(Test_ParallelSeriesUpgradeSync, self).setUp()
+        # Juju Status Object and data
+        # self.juju_status = mock.MagicMock()
+        # self.juju_status.applications.__getitem__.return_value = FAKE_STATUS
+        # self.patch_object(upgrade_utils, "model")
+        # self.model.get_status.return_value = self.juju_status
+
+    def test_get_leader_and_non_leaders(self):
+        expected = ({
+            'app/0': {
+                'leader': True,
+                'machine': '0',
+                'subordinates': {
+                    'app-hacluster/0': {
+                        'charm': 'local:trusty/hacluster-0',
+                        'leader': True}}}}, {
+            'app/1': {
+                'machine': '1',
+                'subordinates': {
+                    'app-hacluster/1': {
+                        'charm': 'local:trusty/hacluster-0'}}},
+            'app/2': {
+                'machine': '2',
+                'subordinates': {
+                    'app-hacluster/2': {
+                        'charm': 'local:trusty/hacluster-0'}}}})
+
+        self.assertEqual(
+            expected,
+            upgrade_utils.get_leader_and_non_leaders(FAKE_STATUS)
+        )
+
+    def test_app_config_openstack_charm(self):
+        expected = {
+            'origin': 'openstack-origin',
+            'pause_non_leader_subordinate': True,
+            'pause_non_leader_primary': True,
+            'post_upgrade_functions': [],
+            'pre_upgrade_functions': [],
+            'post_application_upgrade_functions': [],
+            'follower_first': False, }
+        config = upgrade_utils.app_config('keystone')
+        self.assertEqual(expected, config)
+
+    def test_app_config_mongo(self):
+        expected = {
+            'origin': None,
+            'pause_non_leader_subordinate': True,
+            'pause_non_leader_primary': True,
+            'post_upgrade_functions': [],
+            'pre_upgrade_functions': [],
+            'post_application_upgrade_functions': [],
+            'follower_first': True, }
+        config = upgrade_utils.app_config('mongodb')
+        self.assertEqual(expected, config)
+
+    def test_app_config_ceph(self):
+        expected = {
+            'origin': 'source',
+            'pause_non_leader_subordinate': False,
+            'pause_non_leader_primary': False,
+            'post_upgrade_functions': [],
+            'pre_upgrade_functions': [],
+            'post_application_upgrade_functions': [],
+            'follower_first': False, }
+        config = upgrade_utils.app_config('ceph-mon')
+        self.assertEqual(expected, config)
+
+    def test_app_config_percona(self):
+        expected = {
+            'origin': 'source',
+            'pause_non_leader_subordinate': True,
+            'pause_non_leader_primary': True,
+            'post_upgrade_functions': [],
+            'pre_upgrade_functions': [],
+            'post_application_upgrade_functions': [
+                ('zaza.openstack.charm_tests.mysql.utils.'
+                 'complete_cluster_series_upgrade')
+            ],
+            'follower_first': False, }
+        config = upgrade_utils.app_config('percona-cluster')
+        self.assertEqual(expected, config)
+
+
+class AioTestCase(ut_utils.BaseTestCase):
+    def __init__(self, methodName='runTest', loop=None):
+        self.loop = loop or asyncio.get_event_loop()
+        self._function_cache = {}
+        super(AioTestCase, self).__init__(methodName=methodName)
+
+    def coroutine_function_decorator(self, func):
+        def wrapper(*args, **kw):
+            return self.loop.run_until_complete(func(*args, **kw))
+        return wrapper
+
+    def __getattribute__(self, item):
+        attr = object.__getattribute__(self, item)
+        if asyncio.iscoroutinefunction(attr) and item.startswith('test_'):
+            if item not in self._function_cache:
+                self._function_cache[item] = (
+                    self.coroutine_function_decorator(attr))
+            return self._function_cache[item]
+        return attr
+
+
+class TestParallelSeriesUpgrade(AioTestCase):
+    def setUp(self):
+        super(TestParallelSeriesUpgrade, self).setUp()
+        self.patch_object(series_upgrade, "async_prepare_series_upgrade")
+        self.patch_object(generic_utils, 'check_call')
+        # Juju Status Object and data
+
+        self.juju_status = mock.AsyncMock()
+        self.juju_status.return_value.applications.__getitem__.return_value = \
+            FAKE_STATUS
+        self.patch_object(upgrade_utils, "model")
+        self.model.async_get_status = self.juju_status
+        self.async_run_action = mock.AsyncMock()
+        self.model.async_run_action = self.async_run_action
+
+    @mock.patch.object(upgrade_utils.os_utils, 'async_set_origin')
+    @mock.patch.object(upgrade_utils, 'run_post_application_upgrade_functions')
+    @mock.patch.object(
+        upgrade_utils.series_upgrade_utils, 'async_prepare_series_upgrade')
+    @mock.patch.object(upgrade_utils.series_upgrade_utils, 'async_set_series')
+    @mock.patch.object(upgrade_utils, 'maybe_pause_things')
+    @mock.patch.object(upgrade_utils, 'series_upgrade_machine')
+    async def test_parallel_series_upgrade(
+        self,
+        mock_series_upgrade_machine,
+        mock_maybe_pause_things,
+        mock_async_set_series,
+        mock_async_prepare_series_upgrade,
+        mock_post_application_upgrade_functions,
+        mock_async_set_origin,
+    ):
+        await upgrade_utils.parallel_series_upgrade(
+            'app',
+            from_series='trusty',
+            to_series='xenial',
+        )
+        mock_async_set_series.assert_called_once_with(
+            'app', to_series='xenial')
+        self.juju_status.assert_called()
+        mock_async_prepare_series_upgrade.assert_has_calls([
+            mock.call('1', to_series='xenial'),
+            mock.call('2', to_series='xenial'),
+            mock.call('0', to_series='xenial'),
+        ])
+        mock_maybe_pause_things.assert_called()
+        mock_series_upgrade_machine.assert_has_calls([
+            mock.call(
+                '1',
+                files=None,
+                workaround_script=None,
+                post_upgrade_functions=None),
+            mock.call(
+                '2',
+                files=None,
+                workaround_script=None,
+                post_upgrade_functions=None),
+            mock.call(
+                '0',
+                files=None,
+                workaround_script=None,
+                post_upgrade_functions=None),
+        ])
+        mock_async_set_origin.assert_called_once_with(
+            'app', 'openstack-origin')
+        mock_post_application_upgrade_functions.assert_called_once_with(None)
+
+    @mock.patch.object(upgrade_utils.os_utils, 'async_set_origin')
+    @mock.patch.object(upgrade_utils, 'run_post_application_upgrade_functions')
+    @mock.patch.object(
+        upgrade_utils.series_upgrade_utils, 'async_prepare_series_upgrade')
+    @mock.patch.object(upgrade_utils.series_upgrade_utils, 'async_set_series')
+    @mock.patch.object(upgrade_utils, 'maybe_pause_things')
+    @mock.patch.object(upgrade_utils, 'series_upgrade_machine')
+    async def test_serial_series_upgrade(
+        self,
+        mock_series_upgrade_machine,
+        mock_maybe_pause_things,
+        mock_async_set_series,
+        mock_async_prepare_series_upgrade,
+        mock_post_application_upgrade_functions,
+        mock_async_set_origin,
+    ):
+        await upgrade_utils.serial_series_upgrade(
+            'app',
+            from_series='trusty',
+            to_series='xenial',
+        )
+        mock_async_set_series.assert_called_once_with(
+            'app', to_series='xenial')
+        self.juju_status.assert_called()
+        mock_async_prepare_series_upgrade.assert_has_calls([
+            mock.call('0', to_series='xenial'),
+            mock.call('1', to_series='xenial'),
+            mock.call('2', to_series='xenial'),
+        ])
+        mock_maybe_pause_things.assert_called()
+        mock_series_upgrade_machine.assert_has_calls([
+            mock.call(
+                '0',
+                files=None,
+                workaround_script=None,
+                post_upgrade_functions=None),
+            mock.call(
+                '1',
+                files=None,
+                workaround_script=None,
+                post_upgrade_functions=None),
+            mock.call(
+                '2',
+                files=None,
+                workaround_script=None,
+                post_upgrade_functions=None),
+        ])
+        mock_async_set_origin.assert_called_once_with(
+            'app', 'openstack-origin')
+        mock_post_application_upgrade_functions.assert_called_once_with(None)
+
+    @mock.patch.object(
+        upgrade_utils.series_upgrade_utils, 'async_complete_series_upgrade')
+    @mock.patch.object(upgrade_utils, 'reboot')
+    @mock.patch.object(upgrade_utils, 'async_do_release_upgrade')
+    @mock.patch.object(upgrade_utils, 'async_dist_upgrade')
+    async def test_series_upgrade_machine(
+        self,
+        mock_async_dist_upgrade,
+        mock_async_do_release_upgrade,
+        mock_reboot,
+        mock_async_complete_series_upgrade
+    ):
+        await upgrade_utils.series_upgrade_machine(
+            '1',
+            post_upgrade_functions=None,
+            pre_upgrade_functions=None,
+            files=None,
+            workaround_script=None)
+        mock_async_dist_upgrade.assert_called_once_with('1')
+        mock_async_do_release_upgrade.assert_called_once_with('1')
+        mock_reboot.assert_called_once_with('1')
+        mock_async_complete_series_upgrade.assert_called_once_with('1')
+
+    async def test_maybe_pause_things_primary(self):
+        await upgrade_utils.maybe_pause_things(
+            FAKE_STATUS,
+            ['app/1', 'app/2'],
+            pause_non_leader_subordinate=False,
+            pause_non_leader_primary=True)
+        self.async_run_action.assert_has_calls([
+            mock.call('app/1', "pause", action_params={}),
+            mock.call('app/2', "pause", action_params={}),
+        ])
+
+    async def test_maybe_pause_things_subordinates(self):
+        await upgrade_utils.maybe_pause_things(
+            FAKE_STATUS,
+            ['app/1', 'app/2'],
+            pause_non_leader_subordinate=True,
+            pause_non_leader_primary=False)
+        self.async_run_action.assert_has_calls([
+            mock.call('app-hacluster/1', "pause", action_params={}),
+            mock.call('app-hacluster/2', "pause", action_params={}),
+        ])
+
+    async def test_maybe_pause_things_all(self):
+        await upgrade_utils.maybe_pause_things(
+            FAKE_STATUS,
+            ['app/1', 'app/2'],
+            pause_non_leader_subordinate=True,
+            pause_non_leader_primary=True)
+        self.async_run_action.assert_has_calls([
+            mock.call('app-hacluster/1', "pause", action_params={}),
+            mock.call('app/1', "pause", action_params={}),
+            mock.call('app-hacluster/2', "pause", action_params={}),
+            mock.call('app/2', "pause", action_params={}),
+        ])
+
+    async def test_maybe_pause_things_none(self):
+        await upgrade_utils.maybe_pause_things(
+            FAKE_STATUS,
+            ['app/1', 'app/2'],
+            pause_non_leader_subordinate=False,
+            pause_non_leader_primary=False)
+        self.async_run_action.assert_not_called()
+
+    @mock.patch.object(upgrade_utils, 'run_on_machine')
+    async def test_async_do_release_upgrade(self, mock_run_on_machine):
+        await upgrade_utils.async_do_release_upgrade('1')
+        do_release_upgrade_cmd = (
+            'yes | sudo DEBIAN_FRONTEND=noninteractive '
+            'do-release-upgrade -f DistUpgradeViewNonInteractive')
+        mock_run_on_machine.assert_called_once_with(
+            '1', do_release_upgrade_cmd, timeout='120m'
+        )
+
+    async def test_prepare_series_upgrade(self):
+        await upgrade_utils.prepare_series_upgrade(
+            '1', to_series='xenial'
+        )
+        self.async_prepare_series_upgrade.assert_called_once_with(
+            '1', to_series='xenial'
+        )
+
+    @mock.patch.object(upgrade_utils, 'run_on_machine')
+    async def test_reboot(self, mock_run_on_machine):
+        await upgrade_utils.reboot('1')
+        mock_run_on_machine.assert_called_once_with(
+            '1', 'shutdown --reboot now & exit'
+        )
+
+    async def test_run_on_machine(self):
+        await upgrade_utils.run_on_machine('1', 'test')
+        self.check_call.assert_called_once_with(
+            ['juju', 'run', '--machine=1', 'test'])
+
+    async def test_run_on_machine_with_timeout(self):
+        await upgrade_utils.run_on_machine('1', 'test', timeout='20m')
+        self.check_call.assert_called_once_with(
+            ['juju', 'run', '--machine=1', '--timeout=20m', 'test'])
+
+    async def test_run_on_machine_with_model(self):
+        await upgrade_utils.run_on_machine('1', 'test', model_name='test')
+        self.check_call.assert_called_once_with(
+            ['juju', 'run', '--machine=1', '--model=test', 'test'])
+
+    @mock.patch.object(upgrade_utils, 'run_on_machine')
+    async def test_async_dist_upgrade(self, mock_run_on_machine):
+        await upgrade_utils.async_dist_upgrade('1')
+        apt_update_command = (
+            """yes | sudo DEBIAN_FRONTEND=noninteractive apt --assume-yes """
+            """-o "Dpkg::Options::=--force-confdef" """
+            """-o "Dpkg::Options::=--force-confold" dist-upgrade""")
+        mock_run_on_machine.assert_has_calls([
+            mock.call('1', 'sudo apt update'),
+            mock.call('1', apt_update_command),
+        ])
diff --git a/zaza/openstack/utilities/parallel_series_upgrade.py b/zaza/openstack/utilities/parallel_series_upgrade.py
index 87deb51..7e6693e 100755
--- a/zaza/openstack/utilities/parallel_series_upgrade.py
+++ b/zaza/openstack/utilities/parallel_series_upgrade.py
@@ -187,7 +187,9 @@ async def parallel_series_upgrade(
         .format(application, follower_first))
 
     status = (await model.async_get_status()).applications[application]
-    leader, non_leaders = await get_leader_and_non_leaders(status, application)
+    logging.info(
+        "Configuring leader / non leaders for {}".format(application))
+    leader, non_leaders = get_leader_and_non_leaders(status)
     for leader_name, leader_unit in leader.items():
         leader_machine = leader_unit["machine"]
         leader = leader_name
@@ -205,10 +207,12 @@ async def parallel_series_upgrade(
         application, to_series=to_series)
 
     prepare_group = [
-        prepare_series_upgrade(machine, to_series=to_series)
+        series_upgrade_utils.async_prepare_series_upgrade(
+            machine, to_series=to_series)
         for machine in machines]
     asyncio.gather(*prepare_group)
-    await prepare_series_upgrade(leader_machine, to_series=to_series)
+    await series_upgrade_utils.async_prepare_series_upgrade(
+        leader_machine, to_series=to_series)
     if leader_machine not in completed_machines:
         machines.append(leader_machine)
     upgrade_group = [
@@ -285,7 +289,9 @@ async def serial_series_upgrade(
         "About to upgrade the units of {} in serial (follower first: {})"
         .format(application, follower_first))
     status = (await model.async_get_status()).applications[application]
-    leader, non_leaders = await get_leader_and_non_leaders(application)
+    logging.info(
+        "Configuring leader / non leaders for {}".format(application))
+    leader, non_leaders = get_leader_and_non_leaders(status)
     for leader_name, leader_unit in leader.items():
         leader_machine = leader_unit["machine"]
         leader = leader_name
@@ -304,7 +310,8 @@ async def serial_series_upgrade(
         application, to_series=to_series)
 
     if not follower_first and leader_machine not in completed_machines:
-        await prepare_series_upgrade(leader_machine, to_series)
+        await series_upgrade_utils.async_prepare_series_upgrade(
+            leader_machine, to_series=to_series)
         logging.info("About to upgrade leader of {}: {}"
                      .format(application, leader_machine))
         await series_upgrade_machine(
@@ -313,7 +320,8 @@ async def serial_series_upgrade(
             post_upgrade_functions=post_upgrade_functions)
 
     for machine in machines:
-        await prepare_series_upgrade(machine, to_series)
+        await series_upgrade_utils.async_prepare_series_upgrade(
+            machine, to_series=to_series)
         logging.info("About to upgrade follower of {}: {}"
                      .format(application, machine))
         await series_upgrade_machine(
@@ -322,7 +330,8 @@ async def serial_series_upgrade(
             post_upgrade_functions=post_upgrade_functions)
 
     if follower_first and leader_machine not in completed_machines:
-        await prepare_series_upgrade(leader_machine, to_series)
+        await series_upgrade_utils.async_prepare_series_upgrade(
+            leader_machine, to_series=to_series)
         logging.info("About to upgrade leader of {}: {}"
                      .format(application, leader_machine))
         await series_upgrade_machine(
@@ -360,7 +369,7 @@ async def series_upgrade_machine(
     logging.info(
         "About to dist-upgrade ({})".format(machine))
 
-    run_pre_upgrade_functions(pre_upgrade_functions)
+    run_pre_upgrade_functions(machine, pre_upgrade_functions)
     await async_dist_upgrade(machine)
     await async_do_release_upgrade(machine)
     await reboot(machine)
@@ -436,7 +445,7 @@ async def maybe_pause_things(
     await asyncio.gather(*leader_pauses)
 
 
-async def get_leader_and_non_leaders(status, application):
+def get_leader_and_non_leaders(status):
     """Get the leader and non-leader Juju units.
 
     This function returns a tuple that looks like:
@@ -452,13 +461,9 @@ async def get_leader_and_non_leaders(status, application):
     The first entry of this tuple is the leader, and the second is
     all non-leader units.
 
-    :param application: Application to fetch details for
-    :type application: str
     :returns: A tuple of dicts identifying leader and non-leaders
     :rtype: Dict[str, List[juju.Unit]]
     """
-    logging.info(
-        "Configuring leader / non leaders for {}".format(application))
     leader = None
     non_leaders = {}
     for name, unit in status["units"].items():
@@ -487,42 +492,22 @@ async def prepare_series_upgrade(machine, to_series):
         machine, to_series=to_series)
 
 
-async def reboot(unit):
+async def reboot(machine):
     """Reboot the named machine.
 
-    :param unit: Machine to reboot
-    :type unit: str
+    :param machine: Machine to reboot
+    :type machine: str
     :returns: Nothing
     :rtype: None
     """
     try:
-        await run_on_machine(unit, 'shutdown --reboot now & exit')
+        await run_on_machine(machine, 'shutdown --reboot now & exit')
         # await run_on_machine(unit, "sudo reboot && exit")
     except subprocess.CalledProcessError as e:
         logging.warn("Error doing reboot: {}".format(e))
         pass
 
 
-async def complete_series_upgrade(machines, to_series):
-    """Execute juju series-upgrade complete on machine.
-
-    NOTE: This is a new feature in juju behind a feature flag and not yet in
-    libjuju.
-    export JUJU_DEV_FEATURE_FLAGS=upgrade-series
-    :param machine_num: Machine number
-    :type machine_num: str
-    :returns: None
-    :rtype: None
-    """
-    logging.info("Completing series upgrade for {}".format(machines))
-    group = []
-    for machine in machines:
-        # This can fail on the non-leaders if the leader goes first?
-        group.append(
-            series_upgrade_utils.async_complete_series_upgrade(machine))
-    await asyncio.gather(*group)
-
-
 async def run_on_machine(machine, command, model_name=None, timeout=None):
     """Juju run on unit.
 
diff --git a/zaza/openstack/utilities/upgrade_utils.py b/zaza/openstack/utilities/upgrade_utils.py
index c864b9b..4066d19 100644
--- a/zaza/openstack/utilities/upgrade_utils.py
+++ b/zaza/openstack/utilities/upgrade_utils.py
@@ -20,8 +20,8 @@ import zaza.model
 
 
 SERVICE_GROUPS = collections.OrderedDict([
-    ('Stateful Services', ['percona-cluster', 'rabbitmq-server']),
-    ('Core Identity', ['keystone', 'ceph-mon']),
+    ('Stateful Services', ['percona-cluster', 'rabbitmq-server', 'ceph-mon']),
+    ('Core Identity', ['keystone']),
     ('Control Plane', [
         'aodh', 'barbican', 'ceilometer', 'ceph-fs',
         'ceph-radosgw', 'cinder', 'designate',

From 7649bcb10e364374ea6537a5cd7b2802aac5927a Mon Sep 17 00:00:00 2001
From: Chris MacNaughton <chris.macnaughton@canonical.com>
Date: Fri, 10 Apr 2020 09:29:07 +0200
Subject: [PATCH 07/19] Skip the async unit tests on py35

---
 .../utilities/test_zaza_utilities_parallel_series_upgrade.py  | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/unit_tests/utilities/test_zaza_utilities_parallel_series_upgrade.py b/unit_tests/utilities/test_zaza_utilities_parallel_series_upgrade.py
index 966b54e..3935175 100644
--- a/unit_tests/utilities/test_zaza_utilities_parallel_series_upgrade.py
+++ b/unit_tests/utilities/test_zaza_utilities_parallel_series_upgrade.py
@@ -14,6 +14,8 @@
 
 import asyncio
 import mock
+import sys
+import unittest
 import unit_tests.utils as ut_utils
 import zaza.openstack.utilities.generic as generic_utils
 import zaza.openstack.utilities.series_upgrade as series_upgrade
@@ -149,6 +151,8 @@ class AioTestCase(ut_utils.BaseTestCase):
 class TestParallelSeriesUpgrade(AioTestCase):
     def setUp(self):
         super(TestParallelSeriesUpgrade, self).setUp()
+        if sys.version_info < (3, 6, 0):
+            raise unittest.SkipTest("Can't AsyncMock in py35")
         self.patch_object(series_upgrade, "async_prepare_series_upgrade")
         self.patch_object(generic_utils, 'check_call')
         # Juju Status Object and data

From e2d6a5c0810e6d0ba653816230e06e81cf1cb606 Mon Sep 17 00:00:00 2001
From: Chris MacNaughton <chris.macnaughton@canonical.com>
Date: Fri, 10 Apr 2020 09:58:56 +0200
Subject: [PATCH 08/19] Add test coverage for follower-first upgrades

---
 ..._zaza_utilities_parallel_series_upgrade.py | 120 ++++++++++++++++++
 1 file changed, 120 insertions(+)

diff --git a/unit_tests/utilities/test_zaza_utilities_parallel_series_upgrade.py b/unit_tests/utilities/test_zaza_utilities_parallel_series_upgrade.py
index 3935175..7ec7df6 100644
--- a/unit_tests/utilities/test_zaza_utilities_parallel_series_upgrade.py
+++ b/unit_tests/utilities/test_zaza_utilities_parallel_series_upgrade.py
@@ -40,6 +40,18 @@ FAKE_STATUS = {
                             'app-hacluster/2': {
                                    'charm': 'local:trusty/hacluster-0'}}}}}
 
+FAKE_STATUS_MONGO = {
+    'can-upgrade-to': '',
+    'charm': 'local:trusty/mongodb-10',
+    'subordinate-to': [],
+    'units': {'mongo/0': {'leader': True,
+                          'machine': '0',
+                          'subordinates': {}},
+              'mongo/1': {'machine': '1',
+                          'subordinates': {}},
+              'mongo/2': {'machine': '2',
+                          'subordinates': {}}}}
+
 
 class Test_ParallelSeriesUpgradeSync(ut_utils.BaseTestCase):
     def setUp(self):
@@ -165,6 +177,114 @@ class TestParallelSeriesUpgrade(AioTestCase):
         self.async_run_action = mock.AsyncMock()
         self.model.async_run_action = self.async_run_action
 
+    @mock.patch.object(upgrade_utils.os_utils, 'async_set_origin')
+    @mock.patch.object(upgrade_utils, 'run_post_application_upgrade_functions')
+    @mock.patch.object(
+        upgrade_utils.series_upgrade_utils, 'async_prepare_series_upgrade')
+    @mock.patch.object(upgrade_utils.series_upgrade_utils, 'async_set_series')
+    @mock.patch.object(upgrade_utils, 'maybe_pause_things')
+    @mock.patch.object(upgrade_utils, 'series_upgrade_machine')
+    async def test_parallel_series_upgrade_mongo(
+        self,
+        mock_series_upgrade_machine,
+        mock_maybe_pause_things,
+        mock_async_set_series,
+        mock_async_prepare_series_upgrade,
+        mock_post_application_upgrade_functions,
+        mock_async_set_origin,
+    ):
+        self.juju_status.return_value.applications.__getitem__.return_value = \
+            FAKE_STATUS_MONGO
+        upgrade_config = upgrade_utils.app_config('mongodb')
+        await upgrade_utils.parallel_series_upgrade(
+            'mongodb',
+            from_series='trusty',
+            to_series='xenial',
+            **upgrade_config
+        )
+        mock_async_set_series.assert_called_once_with(
+            'mongodb', to_series='xenial')
+        self.juju_status.assert_called()
+        mock_async_prepare_series_upgrade.assert_has_calls([
+            mock.call('1', to_series='xenial'),
+            mock.call('2', to_series='xenial'),
+            mock.call('0', to_series='xenial'),
+        ])
+        mock_maybe_pause_things.assert_called()
+        mock_series_upgrade_machine.assert_has_calls([
+            mock.call(
+                '1',
+                files=None,
+                workaround_script=None,
+                post_upgrade_functions=[]),
+            mock.call(
+                '2',
+                files=None,
+                workaround_script=None,
+                post_upgrade_functions=[]),
+            mock.call(
+                '0',
+                files=None,
+                workaround_script=None,
+                post_upgrade_functions=[]),
+        ])
+        mock_async_set_origin.assert_not_called()
+        mock_post_application_upgrade_functions.assert_called_once_with([])
+
+    @mock.patch.object(upgrade_utils.os_utils, 'async_set_origin')
+    @mock.patch.object(upgrade_utils, 'run_post_application_upgrade_functions')
+    @mock.patch.object(
+        upgrade_utils.series_upgrade_utils, 'async_prepare_series_upgrade')
+    @mock.patch.object(upgrade_utils.series_upgrade_utils, 'async_set_series')
+    @mock.patch.object(upgrade_utils, 'maybe_pause_things')
+    @mock.patch.object(upgrade_utils, 'series_upgrade_machine')
+    async def test_serial_series_upgrade_mongo(
+        self,
+        mock_series_upgrade_machine,
+        mock_maybe_pause_things,
+        mock_async_set_series,
+        mock_async_prepare_series_upgrade,
+        mock_post_application_upgrade_functions,
+        mock_async_set_origin,
+    ):
+        self.juju_status.return_value.applications.__getitem__.return_value = \
+            FAKE_STATUS_MONGO
+        upgrade_config = upgrade_utils.app_config('mongodb')
+        await upgrade_utils.serial_series_upgrade(
+            'mongodb',
+            from_series='trusty',
+            to_series='xenial',
+            **upgrade_config
+        )
+        mock_async_set_series.assert_called_once_with(
+            'mongodb', to_series='xenial')
+        self.juju_status.assert_called()
+        mock_async_prepare_series_upgrade.assert_has_calls([
+            mock.call('1', to_series='xenial'),
+            mock.call('2', to_series='xenial'),
+            mock.call('0', to_series='xenial'),
+        ])
+        mock_maybe_pause_things.assert_called()
+        mock_series_upgrade_machine.assert_has_calls([
+            mock.call(
+                '1',
+                files=None,
+                workaround_script=None,
+                post_upgrade_functions=[]),
+            mock.call(
+                '2',
+                files=None,
+                workaround_script=None,
+                post_upgrade_functions=[]),
+            mock.call(
+                '0',
+                files=None,
+                workaround_script=None,
+                post_upgrade_functions=[]),
+        ])
+        mock_async_set_origin.assert_not_called()
+        mock_post_application_upgrade_functions.assert_called_once_with([])
+
     @mock.patch.object(upgrade_utils.os_utils, 'async_set_origin')
     @mock.patch.object(upgrade_utils, 'run_post_application_upgrade_functions')
     @mock.patch.object(

From 4d38a158f0fd070390e5872697aeb6e8ba5aa3e8 Mon Sep 17 00:00:00 2001
From: Chris MacNaughton <chris.macnaughton@canonical.com>
Date: Fri, 10 Apr 2020 11:02:22 +0200
Subject: [PATCH 09/19] Tox coverage reporting shoulg be zaza.openstack only

---
 tox.ini | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tox.ini b/tox.ini
index 6a37cfe..7d67e9f 100644
--- a/tox.ini
+++ b/tox.ini
@@ -8,7 +8,7 @@ setenv = VIRTUAL_ENV={envdir}
 install_command =
   pip install {opts} {packages}
 
-commands = nosetests --with-coverage --cover-package=zaza {posargs} {toxinidir}/unit_tests
+commands = nosetests --with-coverage --cover-package=zaza.openstack {posargs} {toxinidir}/unit_tests
 
 [testenv:py3]
 basepython = python3

From 183b3c012fec36a6c9d183b2d7d8dba27f23f28a Mon Sep 17 00:00:00 2001
From: Chris MacNaughton <chris.macnaughton@canonical.com>
Date: Fri, 10 Apr 2020 11:02:38 +0200
Subject: [PATCH 10/19] Add unit tests for new rabbit and mysql helpers

---
 unit_tests/charm_tests/__init__.py            | 13 ++++++++
 unit_tests/charm_tests/test_mysql.py          | 32 +++++++++++++++++++
 .../charm_tests/test_rabbitmq_server.py       | 32 +++++++++++++++++++
 3 files changed, 77 insertions(+)
 create mode 100644 unit_tests/charm_tests/__init__.py
 create mode 100644 unit_tests/charm_tests/test_mysql.py
 create mode 100644 unit_tests/charm_tests/test_rabbitmq_server.py

diff --git a/unit_tests/charm_tests/__init__.py b/unit_tests/charm_tests/__init__.py
new file mode 100644
index 0000000..6131624
--- /dev/null
+++ b/unit_tests/charm_tests/__init__.py
@@ -0,0 +1,13 @@
+# Copyright 2020 Canonical Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
diff --git a/unit_tests/charm_tests/test_mysql.py b/unit_tests/charm_tests/test_mysql.py
new file mode 100644
index 0000000..3493ce9
--- /dev/null
+++ b/unit_tests/charm_tests/test_mysql.py
@@ -0,0 +1,32 @@
+# Copyright 2020 Canonical Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import mock
+import unittest
+
+import zaza.openstack.charm_tests.mysql.utils as mysql_utils
+
+
+class TestMysqlUtils(unittest.TestCase):
+    """Test class to encapsulate testing Mysql test utils."""
+
+    @mock.patch.object(mysql_utils, 'model')
+    def test_mysql_complete_cluster_series_upgrade(self, mock_model):
+        run_action_on_leader = mock.MagicMock()
+        mock_model.run_action_on_leader = run_action_on_leader
+        mysql_utils.complete_cluster_series_upgrade()
+        run_action_on_leader.assert_called_once_with(
+            'mysql',
+            'complete-cluster-series-upgrade',
+            action_params={})
diff --git a/unit_tests/charm_tests/test_rabbitmq_server.py b/unit_tests/charm_tests/test_rabbitmq_server.py
new file mode 100644
index 0000000..8e6ae0d
--- /dev/null
+++ b/unit_tests/charm_tests/test_rabbitmq_server.py
@@ -0,0 +1,32 @@
+# Copyright 2020 Canonical Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import mock
+import unittest
+
+import zaza.openstack.charm_tests.rabbitmq_server.utils as rabbit_utils
+
+
+class TestRabbitUtils(unittest.TestCase):
+    """Test class to encapsulate testing Mysql test utils."""
+
+    @mock.patch.object(rabbit_utils.zaza, 'model')
+    def test_rabbit_complete_cluster_series_upgrade(self, mock_model):
+        run_action_on_leader = mock.MagicMock()
+        mock_model.run_action_on_leader = run_action_on_leader
+        rabbit_utils.complete_cluster_series_upgrade()
+        run_action_on_leader.assert_called_once_with(
+            'rabbitmq-server',
+            'complete-cluster-series-upgrade',
+            action_params={})

From 208439cdf1aec531473518c0957c1c0f8ebb31b3 Mon Sep 17 00:00:00 2001
From: Chris MacNaughton <chris.macnaughton@canonical.com>
Date: Fri, 10 Apr 2020 11:42:21 +0200
Subject: [PATCH 11/19] Add unit tests for pre and post-app upgrade functions

---
 ...est_zaza_utilities_parallel_series_upgrade.py | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/unit_tests/utilities/test_zaza_utilities_parallel_series_upgrade.py b/unit_tests/utilities/test_zaza_utilities_parallel_series_upgrade.py
index 7ec7df6..97371de 100644
--- a/unit_tests/utilities/test_zaza_utilities_parallel_series_upgrade.py
+++ b/unit_tests/utilities/test_zaza_utilities_parallel_series_upgrade.py
@@ -62,6 +62,22 @@ class Test_ParallelSeriesUpgradeSync(ut_utils.BaseTestCase):
         # self.patch_object(upgrade_utils, "model")
         # self.model.get_status.return_value = self.juju_status
 
+    @mock.patch.object(upgrade_utils.cl_utils, 'get_class')
+    def test_run_post_application_upgrade_functions(self, mock_get_class):
+        called = mock.MagicMock()
+        mock_get_class.return_value = called
+        upgrade_utils.run_post_application_upgrade_functions(['my.thing'])
+        mock_get_class.assert_called_once_with('my.thing')
+        called.assert_called()
+
+    @mock.patch.object(upgrade_utils.cl_utils, 'get_class')
+    def test_run_pre_upgrade_functions(self, mock_get_class):
+        called = mock.MagicMock()
+        mock_get_class.return_value = called
+        upgrade_utils.run_pre_upgrade_functions('1', ['my.thing'])
+        mock_get_class.assert_called_once_with('my.thing')
+        called.assert_called_once_with('1')
+
     def test_get_leader_and_non_leaders(self):
         expected = ({
             'app/0': {

From a7df42a9df9845ec7478819dd9ff6517fbe80c5a Mon Sep 17 00:00:00 2001
From: Chris MacNaughton <chris.macnaughton@canonical.com>
Date: Fri, 10 Apr 2020 16:27:03 +0200
Subject: [PATCH 12/19] According to the series upgrade doc, source should
 update early

---
 zaza/openstack/utilities/parallel_series_upgrade.py | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/zaza/openstack/utilities/parallel_series_upgrade.py b/zaza/openstack/utilities/parallel_series_upgrade.py
index 7e6693e..66583c8 100755
--- a/zaza/openstack/utilities/parallel_series_upgrade.py
+++ b/zaza/openstack/utilities/parallel_series_upgrade.py
@@ -205,7 +205,8 @@ async def parallel_series_upgrade(
         pause_non_leader_primary)
     await series_upgrade_utils.async_set_series(
         application, to_series=to_series)
-
+    if origin:
+        await os_utils.async_set_origin(application, origin)
     prepare_group = [
         series_upgrade_utils.async_prepare_series_upgrade(
             machine, to_series=to_series)
@@ -223,8 +224,6 @@ async def parallel_series_upgrade(
         for machine in machines
     ]
     await asyncio.gather(*upgrade_group)
-    if origin:
-        await os_utils.async_set_origin(application, origin)
     run_post_application_upgrade_functions(post_application_upgrade_functions)
 
 
@@ -308,7 +307,8 @@ async def serial_series_upgrade(
         pause_non_leader_primary)
     await series_upgrade_utils.async_set_series(
         application, to_series=to_series)
-
+    if origin:
+        await os_utils.async_set_origin(application, origin)
     if not follower_first and leader_machine not in completed_machines:
         await series_upgrade_utils.async_prepare_series_upgrade(
             leader_machine, to_series=to_series)
@@ -338,8 +338,6 @@ async def serial_series_upgrade(
             leader_machine,
             files=files, workaround_script=workaround_script,
             post_upgrade_functions=post_upgrade_functions)
-    if origin:
-        await os_utils.async_set_origin(application, origin)
     run_post_application_upgrade_functions(post_application_upgrade_functions)
 
 

From 27e5f7d8bc094a498c6000152849e2c365bb4e13 Mon Sep 17 00:00:00 2001
From: Chris MacNaughton <chris.macnaughton@canonical.com>
Date: Fri, 10 Apr 2020 16:51:41 +0200
Subject: [PATCH 13/19] We must have async functions for the before/after
 callables

---
 unit_tests/charm_tests/test_mysql.py          | 14 ++++++--
 .../charm_tests/test_rabbitmq_server.py       | 14 ++++++--
 ..._zaza_utilities_parallel_series_upgrade.py | 36 ++++++++++---------
 zaza/openstack/charm_tests/mysql/utils.py     |  4 +--
 .../charm_tests/rabbitmq_server/utils.py      |  4 +--
 zaza/openstack/charm_tests/vault/setup.py     | 14 ++++++++
 .../utilities/parallel_series_upgrade.py      | 27 ++++++++++----
 7 files changed, 80 insertions(+), 33 deletions(-)

diff --git a/unit_tests/charm_tests/test_mysql.py b/unit_tests/charm_tests/test_mysql.py
index 3493ce9..e38460b 100644
--- a/unit_tests/charm_tests/test_mysql.py
+++ b/unit_tests/charm_tests/test_mysql.py
@@ -12,8 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import asyncio
 import mock
 import unittest
+import sys
 
 import zaza.openstack.charm_tests.mysql.utils as mysql_utils
 
@@ -21,11 +23,17 @@ import zaza.openstack.charm_tests.mysql.utils as mysql_utils
 class TestMysqlUtils(unittest.TestCase):
     """Test class to encapsulate testing Mysql test utils."""
 
+    def setUp(self):
+        super(TestMysqlUtils, self).setUp()
+        if sys.version_info < (3, 6, 0):
+            raise unittest.SkipTest("Can't AsyncMock in py35")
+
     @mock.patch.object(mysql_utils, 'model')
     def test_mysql_complete_cluster_series_upgrade(self, mock_model):
-        run_action_on_leader = mock.MagicMock()
-        mock_model.run_action_on_leader = run_action_on_leader
-        mysql_utils.complete_cluster_series_upgrade()
+        run_action_on_leader = mock.AsyncMock()
+        mock_model.async_run_action_on_leader = run_action_on_leader
+        asyncio.get_event_loop().run_until_complete(
+            mysql_utils.complete_cluster_series_upgrade())
         run_action_on_leader.assert_called_once_with(
             'mysql',
             'complete-cluster-series-upgrade',
diff --git a/unit_tests/charm_tests/test_rabbitmq_server.py b/unit_tests/charm_tests/test_rabbitmq_server.py
index 8e6ae0d..e092122 100644
--- a/unit_tests/charm_tests/test_rabbitmq_server.py
+++ b/unit_tests/charm_tests/test_rabbitmq_server.py
@@ -12,8 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import asyncio
 import mock
 import unittest
+import sys
 
 import zaza.openstack.charm_tests.rabbitmq_server.utils as rabbit_utils
 
@@ -21,11 +23,17 @@ import zaza.openstack.charm_tests.rabbitmq_server.utils as rabbit_utils
 class TestRabbitUtils(unittest.TestCase):
     """Test class to encapsulate testing Mysql test utils."""
 
+    def setUp(self):
+        super(TestRabbitUtils, self).setUp()
+        if sys.version_info < (3, 6, 0):
+            raise unittest.SkipTest("Can't AsyncMock in py35")
+
     @mock.patch.object(rabbit_utils.zaza, 'model')
     def test_rabbit_complete_cluster_series_upgrade(self, mock_model):
-        run_action_on_leader = mock.MagicMock()
-        mock_model.run_action_on_leader = run_action_on_leader
-        rabbit_utils.complete_cluster_series_upgrade()
+        run_action_on_leader = mock.AsyncMock()
+        mock_model.async_run_action_on_leader = run_action_on_leader
+        asyncio.get_event_loop().run_until_complete(
+            rabbit_utils.complete_cluster_series_upgrade())
         run_action_on_leader.assert_called_once_with(
             'rabbitmq-server',
             'complete-cluster-series-upgrade',
diff --git a/unit_tests/utilities/test_zaza_utilities_parallel_series_upgrade.py b/unit_tests/utilities/test_zaza_utilities_parallel_series_upgrade.py
index 97371de..8c17cbb 100644
--- a/unit_tests/utilities/test_zaza_utilities_parallel_series_upgrade.py
+++ b/unit_tests/utilities/test_zaza_utilities_parallel_series_upgrade.py
@@ -62,22 +62,6 @@ class Test_ParallelSeriesUpgradeSync(ut_utils.BaseTestCase):
         # self.patch_object(upgrade_utils, "model")
         # self.model.get_status.return_value = self.juju_status
 
-    @mock.patch.object(upgrade_utils.cl_utils, 'get_class')
-    def test_run_post_application_upgrade_functions(self, mock_get_class):
-        called = mock.MagicMock()
-        mock_get_class.return_value = called
-        upgrade_utils.run_post_application_upgrade_functions(['my.thing'])
-        mock_get_class.assert_called_once_with('my.thing')
-        called.assert_called()
-
-    @mock.patch.object(upgrade_utils.cl_utils, 'get_class')
-    def test_run_pre_upgrade_functions(self, mock_get_class):
-        called = mock.MagicMock()
-        mock_get_class.return_value = called
-        upgrade_utils.run_pre_upgrade_functions('1', ['my.thing'])
-        mock_get_class.assert_called_once_with('my.thing')
-        called.assert_called_once_with('1')
-
     def test_get_leader_and_non_leaders(self):
         expected = ({
             'app/0': {
@@ -193,6 +177,26 @@ class TestParallelSeriesUpgrade(AioTestCase):
         self.async_run_action = mock.AsyncMock()
         self.model.async_run_action = self.async_run_action
 
+    @mock.patch.object(upgrade_utils.cl_utils, 'get_class')
+    async def test_run_post_application_upgrade_functions(
+        self,
+        mock_get_class
+    ):
+        called = mock.AsyncMock()
+        mock_get_class.return_value = called
+        await upgrade_utils.run_post_application_upgrade_functions(
+            ['my.thing'])
+        mock_get_class.assert_called_once_with('my.thing')
+        called.assert_called()
+
+    @mock.patch.object(upgrade_utils.cl_utils, 'get_class')
+    async def test_run_pre_upgrade_functions(self, mock_get_class):
+        called = mock.AsyncMock()
+        mock_get_class.return_value = called
+        await upgrade_utils.run_pre_upgrade_functions('1', ['my.thing'])
+        mock_get_class.assert_called_once_with('my.thing')
+        called.assert_called_once_with('1')
+
     @mock.patch.object(upgrade_utils.os_utils, 'async_set_origin')
     @mock.patch.object(upgrade_utils, 'run_post_application_upgrade_functions')
     @mock.patch.object(
diff --git a/zaza/openstack/charm_tests/mysql/utils.py b/zaza/openstack/charm_tests/mysql/utils.py
index 2622903..1fe5114 100644
--- a/zaza/openstack/charm_tests/mysql/utils.py
+++ b/zaza/openstack/charm_tests/mysql/utils.py
@@ -17,10 +17,10 @@
 import zaza.model as model
 
 
-def complete_cluster_series_upgrade():
+async def complete_cluster_series_upgrade():
     """Run the complete-cluster-series-upgrade action on the lead unit."""
     # TODO: Make this work across either mysql or percona-cluster names
-    model.run_action_on_leader(
+    await model.async_run_action_on_leader(
         'mysql',
         'complete-cluster-series-upgrade',
         action_params={})
diff --git a/zaza/openstack/charm_tests/rabbitmq_server/utils.py b/zaza/openstack/charm_tests/rabbitmq_server/utils.py
index 62db7e6..3994178 100644
--- a/zaza/openstack/charm_tests/rabbitmq_server/utils.py
+++ b/zaza/openstack/charm_tests/rabbitmq_server/utils.py
@@ -586,9 +586,9 @@ def _post_check_unit_cluster_nodes(unit, nodes, unit_node_names):
     return errors
 
 
-def complete_cluster_series_upgrade():
+async def complete_cluster_series_upgrade():
     """Run the complete-cluster-series-upgrade action on the lead unit."""
-    zaza.model.run_action_on_leader(
+    await zaza.model.async_run_action_on_leader(
         'rabbitmq-server',
         'complete-cluster-series-upgrade',
         action_params={})
diff --git a/zaza/openstack/charm_tests/vault/setup.py b/zaza/openstack/charm_tests/vault/setup.py
index 96c45cb..21c3793 100644
--- a/zaza/openstack/charm_tests/vault/setup.py
+++ b/zaza/openstack/charm_tests/vault/setup.py
@@ -66,6 +66,20 @@ def mojo_unseal_by_unit():
             zaza.model.run_on_unit(unit_name, './hooks/update-status')
 
 
+async def async_mojo_unseal_by_unit():
+    """Unseal any units reported as sealed using mojo cacert."""
+    cacert = zaza.openstack.utilities.generic.get_mojo_cacert_path()
+    vault_creds = vault_utils.get_credentails()
+    for client in vault_utils.get_clients(cacert=cacert):
+        if client.hvac_client.is_sealed():
+            client.hvac_client.unseal(vault_creds['keys'][0])
+            unit_name = await juju_utils.async_get_unit_name_from_ip_address(
+                client.addr,
+                'vault')
+            await zaza.model.async_run_on_unit(
+                unit_name, './hooks/update-status')
+
+
 def auto_initialize(cacert=None, validation_application='keystone'):
     """Auto initialize vault for testing.
 
diff --git a/zaza/openstack/utilities/parallel_series_upgrade.py b/zaza/openstack/utilities/parallel_series_upgrade.py
index 66583c8..9e29d6d 100755
--- a/zaza/openstack/utilities/parallel_series_upgrade.py
+++ b/zaza/openstack/utilities/parallel_series_upgrade.py
@@ -93,7 +93,7 @@ def app_config(charm_name):
             'pause_non_leader_subordinate': True,
             'post_upgrade_functions': [
                 ('zaza.openstack.charm_tests.vault.setup.'
-                 'mojo_unseal_by_unit')]
+                 'async_mojo_unseal_by_unit')]
         },
         'mongodb': {
             'origin': None,
@@ -338,7 +338,8 @@ async def serial_series_upgrade(
             leader_machine,
             files=files, workaround_script=workaround_script,
             post_upgrade_functions=post_upgrade_functions)
-    run_post_application_upgrade_functions(post_application_upgrade_functions)
+    await run_post_application_upgrade_functions(
+        post_application_upgrade_functions)
 
 
 async def series_upgrade_machine(
@@ -372,10 +373,10 @@ async def series_upgrade_machine(
     await async_do_release_upgrade(machine)
     await reboot(machine)
     await series_upgrade_utils.async_complete_series_upgrade(machine)
-    series_upgrade_utils.run_post_upgrade_functions(post_upgrade_functions)
+    await run_post_upgrade_functions(post_upgrade_functions)
 
 
-def run_pre_upgrade_functions(machine, pre_upgrade_functions):
+async def run_pre_upgrade_functions(machine, pre_upgrade_functions):
     """Execute list supplied functions.
 
     Each of the supplied functions will be called with a single
@@ -389,10 +390,10 @@ def run_pre_upgrade_functions(machine, pre_upgrade_functions):
     if pre_upgrade_functions:
         for func in pre_upgrade_functions:
             logging.info("Running {}".format(func))
-            cl_utils.get_class(func)(machine)
+            await cl_utils.get_class(func)(machine)
 
 
-def run_post_application_upgrade_functions(post_upgrade_functions):
+async def run_post_upgrade_functions(post_upgrade_functions):
     """Execute list supplied functions.
 
     :param post_upgrade_functions: List of functions
@@ -401,7 +402,19 @@ def run_post_application_upgrade_functions(post_upgrade_functions):
     if post_upgrade_functions:
         for func in post_upgrade_functions:
             logging.info("Running {}".format(func))
-            cl_utils.get_class(func)()
+            await cl_utils.get_class(func)()
+
+
+async def run_post_application_upgrade_functions(post_upgrade_functions):
+    """Execute list supplied functions.
+
+    :param post_upgrade_functions: List of functions
+    :type post_upgrade_functions: [function, function, ...]
+    """
+    if post_upgrade_functions:
+        for func in post_upgrade_functions:
+            logging.info("Running {}".format(func))
+            await cl_utils.get_class(func)()
 
 
 async def maybe_pause_things(

From 6c70ee5171b8193b4f834275215b315e0e288cfe Mon Sep 17 00:00:00 2001
From: Chris MacNaughton <chris.macnaughton@canonical.com>
Date: Tue, 14 Apr 2020 12:42:18 +0200
Subject: [PATCH 14/19] The units in an application need to be idle before we
 try to series-upgrade them

---
 .../utilities/parallel_series_upgrade.py      | 63 ++++++++++++++++++-
 1 file changed, 61 insertions(+), 2 deletions(-)

diff --git a/zaza/openstack/utilities/parallel_series_upgrade.py b/zaza/openstack/utilities/parallel_series_upgrade.py
index 9e29d6d..de740ce 100755
--- a/zaza/openstack/utilities/parallel_series_upgrade.py
+++ b/zaza/openstack/utilities/parallel_series_upgrade.py
@@ -17,6 +17,8 @@
 
 
 import asyncio
+
+import concurrent
 import collections
 import copy
 import logging
@@ -207,11 +209,15 @@ async def parallel_series_upgrade(
         application, to_series=to_series)
     if origin:
         await os_utils.async_set_origin(application, origin)
+    app_idle = [
+        wait_for_unit_idle(unit) for unit in status["units"]
+    ]
+    await asyncio.gather(*app_idle)
     prepare_group = [
         series_upgrade_utils.async_prepare_series_upgrade(
             machine, to_series=to_series)
         for machine in machines]
-    asyncio.gather(*prepare_group)
+    await asyncio.gather(*prepare_group)
     await series_upgrade_utils.async_prepare_series_upgrade(
         leader_machine, to_series=to_series)
     if leader_machine not in completed_machines:
@@ -224,6 +230,7 @@ async def parallel_series_upgrade(
         for machine in machines
     ]
     await asyncio.gather(*upgrade_group)
+    completed_machines.extend(machines)
     run_post_application_upgrade_functions(post_application_upgrade_functions)
 
 
@@ -310,6 +317,7 @@ async def serial_series_upgrade(
     if origin:
         await os_utils.async_set_origin(application, origin)
     if not follower_first and leader_machine not in completed_machines:
+        await wait_for_unit_idle(leader)
         await series_upgrade_utils.async_prepare_series_upgrade(
             leader_machine, to_series=to_series)
         logging.info("About to upgrade leader of {}: {}"
@@ -318,8 +326,14 @@ async def serial_series_upgrade(
             leader_machine,
             files=files, workaround_script=workaround_script,
             post_upgrade_functions=post_upgrade_functions)
+        completed_machines.append(leader_machine)
 
-    for machine in machines:
+    # for machine in machines:
+    for unit_name, unit in non_leaders.items():
+        machine = unit['machine']
+        if machine in completed_machines:
+            continue
+        await wait_for_unit_idle(unit_name)
         await series_upgrade_utils.async_prepare_series_upgrade(
             machine, to_series=to_series)
         logging.info("About to upgrade follower of {}: {}"
@@ -328,8 +342,10 @@ async def serial_series_upgrade(
             machine,
             files=files, workaround_script=workaround_script,
             post_upgrade_functions=post_upgrade_functions)
+        completed_machines.append(machine)
 
     if follower_first and leader_machine not in completed_machines:
+        await wait_for_unit_idle(leader)
         await series_upgrade_utils.async_prepare_series_upgrade(
             leader_machine, to_series=to_series)
         logging.info("About to upgrade leader of {}: {}"
@@ -338,6 +354,7 @@ async def serial_series_upgrade(
             leader_machine,
             files=files, workaround_script=workaround_script,
             post_upgrade_functions=post_upgrade_functions)
+        completed_machines.append(leader_machine)
     await run_post_application_upgrade_functions(
         post_application_upgrade_functions)
 
@@ -577,3 +594,45 @@ async def async_do_release_upgrade(machine):
         'do-release-upgrade -f DistUpgradeViewNonInteractive')
 
     await run_on_machine(machine, do_release_upgrade_cmd, timeout="120m")
+
+
+# TODO: Move these functions into zaza.model
+async def wait_for_unit_idle(unit_name, timeout=600):
+    """Wait until the unit's agent is idle.
+
+    :param unit_name: The unit name of the application, ex: mysql/0
+    :type unit_name: str
+    :param timeout: How long to wait before timing out
+    :type timeout: int
+    :returns: None
+    :rtype: None
+    """
+    app = unit_name.split('/')[0]
+    try:
+        await model.async_block_until(
+            _unit_idle(app, unit_name),
+            timeout=timeout)
+    except concurrent.futures._base.TimeoutError:
+        raise model.ModelTimeout("Zaza has timed out waiting on {} to "
+                                 "reach idle state.".format(unit_name))
+
+
+def _unit_idle(app, unit_name):
+    async def f():
+        x = await get_agent_status(app, unit_name)
+        return x == "idle"
+    return f
+
+
+async def get_agent_status(app, unit_name):
+    """Get the current status of the specified unit.
+
+    :param app: The name of the Juju application, ex: mysql
+    :type app: str
+    :param unit_name: The unit name of the application, ex: mysql/0
+    :type unit_name: str
+    :returns: The agent status, either active / idle, returned by Juju
+    :rtype: str
+    """
+    return (await model.async_get_status()). \
+        applications[app]['units'][unit_name]['agent-status']['status']

From 592622935cae297983932792b643170a0eddfa27 Mon Sep 17 00:00:00 2001
From: Chris MacNaughton <chris.macnaughton@canonical.com>
Date: Wed, 15 Apr 2020 09:40:37 +0200
Subject: [PATCH 15/19] tidy up and add more debugability

---
 ..._zaza_utilities_parallel_series_upgrade.py | 15 ++++--
 zaza/openstack/utilities/generic.py           |  7 +++
 .../utilities/parallel_series_upgrade.py      | 53 ++++++++-----------
 zaza/openstack/utilities/series_upgrade.py    |  2 +
 4 files changed, 43 insertions(+), 34 deletions(-)

diff --git a/unit_tests/utilities/test_zaza_utilities_parallel_series_upgrade.py b/unit_tests/utilities/test_zaza_utilities_parallel_series_upgrade.py
index 8c17cbb..36638b1 100644
--- a/unit_tests/utilities/test_zaza_utilities_parallel_series_upgrade.py
+++ b/unit_tests/utilities/test_zaza_utilities_parallel_series_upgrade.py
@@ -177,6 +177,9 @@ class TestParallelSeriesUpgrade(AioTestCase):
         self.async_run_action = mock.AsyncMock()
         self.model.async_run_action = self.async_run_action
 
+        self.async_block_until = mock.AsyncMock()
+        self.model.async_block_until = self.async_block_until
+
     @mock.patch.object(upgrade_utils.cl_utils, 'get_class')
     async def test_run_post_application_upgrade_functions(
         self,
@@ -225,11 +228,14 @@ class TestParallelSeriesUpgrade(AioTestCase):
         mock_async_set_series.assert_called_once_with(
             'mongodb', to_series='xenial')
         self.juju_status.assert_called()
+
+        # The below is using `any_order=True` because the ordering is
+        # undetermined and differs between python versions
         mock_async_prepare_series_upgrade.assert_has_calls([
             mock.call('1', to_series='xenial'),
             mock.call('2', to_series='xenial'),
             mock.call('0', to_series='xenial'),
-        ])
+        ], any_order=True)
         mock_maybe_pause_things.assert_called()
         mock_series_upgrade_machine.assert_has_calls([
             mock.call(
@@ -495,7 +501,7 @@ class TestParallelSeriesUpgrade(AioTestCase):
     async def test_reboot(self, mock_run_on_machine):
         await upgrade_utils.reboot('1')
         mock_run_on_machine.assert_called_once_with(
-            '1', 'shutdown --reboot now & exit'
+            '1', 'sudo init 6 & exit'
         )
 
     async def test_run_on_machine(self):
@@ -517,10 +523,11 @@ class TestParallelSeriesUpgrade(AioTestCase):
     async def test_async_dist_upgrade(self, mock_run_on_machine):
         await upgrade_utils.async_dist_upgrade('1')
         apt_update_command = (
-            """yes | sudo DEBIAN_FRONTEND=noninteractive apt --assume-yes """
+            """yes | sudo DEBIAN_FRONTEND=noninteractive """
+            """apt-get --assume-yes """
             """-o "Dpkg::Options::=--force-confdef" """
             """-o "Dpkg::Options::=--force-confold" dist-upgrade""")
         mock_run_on_machine.assert_has_calls([
-            mock.call('1', 'sudo apt update'),
+            mock.call('1', 'sudo apt-get update'),
             mock.call('1', apt_update_command),
         ])
diff --git a/zaza/openstack/utilities/generic.py b/zaza/openstack/utilities/generic.py
index 50dcc69..0874cc6 100644
--- a/zaza/openstack/utilities/generic.py
+++ b/zaza/openstack/utilities/generic.py
@@ -385,10 +385,17 @@ async def check_call(cmd):
         stdout=asyncio.subprocess.PIPE,
         stderr=asyncio.subprocess.PIPE)
     stdout, stderr = await proc.communicate()
+    stdout = stdout.decode('utf-8')
+    stderr = stderr.decode('utf-8')
     if proc.returncode != 0:
         logging.warn("STDOUT: {}".format(stdout))
         logging.warn("STDERR: {}".format(stderr))
         raise subprocess.CalledProcessError(proc.returncode, cmd)
+    else:
+        if stderr:
+            logging.info("STDERR: {} ({})".format(stderr, ' '.join(cmd)))
+        if stdout:
+            logging.info("STDOUT: {} ({})".format(stdout, ' '.join(cmd)))
 
 
 def set_dpkg_non_interactive_on_unit(
diff --git a/zaza/openstack/utilities/parallel_series_upgrade.py b/zaza/openstack/utilities/parallel_series_upgrade.py
index de740ce..dfbd36e 100755
--- a/zaza/openstack/utilities/parallel_series_upgrade.py
+++ b/zaza/openstack/utilities/parallel_series_upgrade.py
@@ -214,12 +214,10 @@ async def parallel_series_upgrade(
     ]
     await asyncio.gather(*app_idle)
     prepare_group = [
-        series_upgrade_utils.async_prepare_series_upgrade(
-            machine, to_series=to_series)
+        prepare_series_upgrade(machine, to_series=to_series)
         for machine in machines]
     await asyncio.gather(*prepare_group)
-    await series_upgrade_utils.async_prepare_series_upgrade(
-        leader_machine, to_series=to_series)
+    await prepare_series_upgrade(leader_machine, to_series=to_series)
     if leader_machine not in completed_machines:
         machines.append(leader_machine)
     upgrade_group = [
@@ -231,7 +229,8 @@ async def parallel_series_upgrade(
     ]
     await asyncio.gather(*upgrade_group)
     completed_machines.extend(machines)
-    run_post_application_upgrade_functions(post_application_upgrade_functions)
+    await run_post_application_upgrade_functions(
+        post_application_upgrade_functions)
 
 
 async def serial_series_upgrade(
@@ -302,11 +301,6 @@ async def serial_series_upgrade(
         leader_machine = leader_unit["machine"]
         leader = leader_name
 
-    machines = [
-        unit["machine"] for name, unit
-        in non_leaders.items()
-        if unit['machine'] not in completed_machines]
-
     await maybe_pause_things(
         status,
         non_leaders,
@@ -318,8 +312,7 @@ async def serial_series_upgrade(
         await os_utils.async_set_origin(application, origin)
     if not follower_first and leader_machine not in completed_machines:
         await wait_for_unit_idle(leader)
-        await series_upgrade_utils.async_prepare_series_upgrade(
-            leader_machine, to_series=to_series)
+        await prepare_series_upgrade(leader_machine, to_series=to_series)
         logging.info("About to upgrade leader of {}: {}"
                      .format(application, leader_machine))
         await series_upgrade_machine(
@@ -334,8 +327,7 @@ async def serial_series_upgrade(
         if machine in completed_machines:
             continue
         await wait_for_unit_idle(unit_name)
-        await series_upgrade_utils.async_prepare_series_upgrade(
-            machine, to_series=to_series)
+        await prepare_series_upgrade(machine, to_series=to_series)
         logging.info("About to upgrade follower of {}: {}"
                      .format(application, machine))
         await series_upgrade_machine(
@@ -346,8 +338,7 @@ async def serial_series_upgrade(
 
     if follower_first and leader_machine not in completed_machines:
         await wait_for_unit_idle(leader)
-        await series_upgrade_utils.async_prepare_series_upgrade(
-            leader_machine, to_series=to_series)
+        await prepare_series_upgrade(leader_machine, to_series=to_series)
         logging.info("About to upgrade leader of {}: {}"
                      .format(application, leader_machine))
         await series_upgrade_machine(
@@ -383,9 +374,8 @@ async def series_upgrade_machine(
     :rtype: None
     """
     logging.info(
-        "About to dist-upgrade ({})".format(machine))
-
-    run_pre_upgrade_functions(machine, pre_upgrade_functions)
+        "About to series-upgrade ({})".format(machine))
+    await run_pre_upgrade_functions(machine, pre_upgrade_functions)
     await async_dist_upgrade(machine)
     await async_do_release_upgrade(machine)
     await reboot(machine)
@@ -401,37 +391,40 @@ async def run_pre_upgrade_functions(machine, pre_upgrade_functions):
 
     :param machine: Machine that is about to be upgraded
     :type machine: str
-    :param pre_upgrade_functions: List of functions
+    :param pre_upgrade_functions: List of awaitable functions
     :type pre_upgrade_functions: [function, function, ...]
     """
     if pre_upgrade_functions:
         for func in pre_upgrade_functions:
             logging.info("Running {}".format(func))
-            await cl_utils.get_class(func)(machine)
+            m = cl_utils.get_class(func)
+            await m(machine)
 
 
 async def run_post_upgrade_functions(post_upgrade_functions):
     """Execute list supplied functions.
 
-    :param post_upgrade_functions: List of functions
+    :param post_upgrade_functions: List of awaitable functions
     :type post_upgrade_functions: [function, function, ...]
     """
     if post_upgrade_functions:
         for func in post_upgrade_functions:
             logging.info("Running {}".format(func))
-            await cl_utils.get_class(func)()
+            m = cl_utils.get_class(func)
+            await m()
 
 
 async def run_post_application_upgrade_functions(post_upgrade_functions):
     """Execute list supplied functions.
 
-    :param post_upgrade_functions: List of functions
+    :param post_upgrade_functions: List of awaitable functions
     :type post_upgrade_functions: [function, function, ...]
     """
     if post_upgrade_functions:
         for func in post_upgrade_functions:
             logging.info("Running {}".format(func))
-            await cl_utils.get_class(func)()
+            m = cl_utils.get_class(func)
+            await m()
 
 
 async def maybe_pause_things(
@@ -515,7 +508,7 @@ async def prepare_series_upgrade(machine, to_series):
     :returns: None
     :rtype: None
     """
-    logging.debug("Preparing series upgrade for: {}".format(machine))
+    logging.info("Preparing series upgrade for: {}".format(machine))
     await series_upgrade_utils.async_prepare_series_upgrade(
         machine, to_series=to_series)
 
@@ -529,7 +522,7 @@ async def reboot(machine):
     :rtype: None
     """
     try:
-        await run_on_machine(machine, 'shutdown --reboot now & exit')
+        await run_on_machine(machine, 'sudo init 6 & exit')
         # await run_on_machine(unit, "sudo reboot && exit")
     except subprocess.CalledProcessError as e:
         logging.warn("Error doing reboot: {}".format(e))
@@ -556,7 +549,7 @@ async def run_on_machine(machine, command, model_name=None, timeout=None):
     if timeout:
         cmd.append('--timeout={}'.format(timeout))
     cmd.append(command)
-    logging.debug("About to call '{}'".format(cmd))
+    logging.info("About to call '{}'".format(cmd))
     await os_utils.check_call(cmd)
 
 
@@ -569,12 +562,12 @@ async def async_dist_upgrade(machine):
     :rtype: None
     """
     logging.info('Updating package db ' + machine)
-    update_cmd = 'sudo apt update'
+    update_cmd = 'sudo apt-get update'
     await run_on_machine(machine, update_cmd)
 
     logging.info('Updating existing packages ' + machine)
     dist_upgrade_cmd = (
-        """yes | sudo DEBIAN_FRONTEND=noninteractive apt --assume-yes """
+        """yes | sudo DEBIAN_FRONTEND=noninteractive apt-get --assume-yes """
         """-o "Dpkg::Options::=--force-confdef" """
         """-o "Dpkg::Options::=--force-confold" dist-upgrade""")
     await run_on_machine(machine, dist_upgrade_cmd)
diff --git a/zaza/openstack/utilities/series_upgrade.py b/zaza/openstack/utilities/series_upgrade.py
index 7f3efa3..f7e79a8 100644
--- a/zaza/openstack/utilities/series_upgrade.py
+++ b/zaza/openstack/utilities/series_upgrade.py
@@ -687,6 +687,7 @@ async def async_complete_series_upgrade(machine_num):
     juju_model = await model.async_get_juju_model()
     cmd = ["juju", "upgrade-series", "-m", juju_model,
            machine_num, "complete"]
+    logging.info("About to call '{}'".format(cmd))
     await os_utils.check_call(cmd)
 
 
@@ -704,6 +705,7 @@ async def async_set_series(application, to_series):
     juju_model = await model.async_get_juju_model()
     cmd = ["juju", "set-series", "-m", juju_model,
            application, to_series]
+    logging.info("About to call '{}'".format(cmd))
     await os_utils.check_call(cmd)
 
 

From c492ecdcac3b2724833c347e978de97ea2e626d7 Mon Sep 17 00:00:00 2001
From: Chris MacNaughton <chris.macnaughton@canonical.com>
Date: Wed, 15 Apr 2020 15:48:37 +0200
Subject: [PATCH 16/19] Add functional test that covers Ubuntu Lite in Parallel

---
 .../series_upgrade/parallel_tests.py          | 53 +++++++++++++++++++
 1 file changed, 53 insertions(+)

diff --git a/zaza/openstack/charm_tests/series_upgrade/parallel_tests.py b/zaza/openstack/charm_tests/series_upgrade/parallel_tests.py
index 5307339..6408cf0 100644
--- a/zaza/openstack/charm_tests/series_upgrade/parallel_tests.py
+++ b/zaza/openstack/charm_tests/series_upgrade/parallel_tests.py
@@ -168,6 +168,59 @@ class BionicFocalSeriesUpgrade(OpenStackParallelSeriesUpgrade):
         cls.to_series = "focal"
 
 
+class UbuntuLiteParallelSeriesUpgrade(unittest.TestCase):
+    """ubuntu Lite Parallel Series Upgrade."""
+
+    @classmethod
+    def setUpClass(cls):
+        """Run setup for Series Upgrades."""
+        cli_utils.setup_logging()
+        cls.from_series = None
+        cls.to_series = None
+
+    def test_200_run_series_upgrade(self):
+        """Run series upgrade."""
+        # Set Feature Flag
+        os.environ["JUJU_DEV_FEATURE_FLAGS"] = "upgrade-series"
+        parallel_series_upgrade.upgrade_ubuntu_lite(
+            from_series=self.from_series,
+            to_series=self.to_series
+        )
+
+
+class TrustyXenialSeriesUpgradeUbuntu(UbuntuLiteParallelSeriesUpgrade):
+    """OpenStack Trusty to Xenial Series Upgrade."""
+
+    @classmethod
+    def setUpClass(cls):
+        """Run setup for Trusty to Xenial Series Upgrades."""
+        super(TrustyXenialSeriesUpgradeUbuntu, cls).setUpClass()
+        cls.from_series = "trusty"
+        cls.to_series = "xenial"
+
+
+class XenialBionicSeriesUpgradeUbuntu(UbuntuLiteParallelSeriesUpgrade):
+    """OpenStack Xenial to Bionic Series Upgrade."""
+
+    @classmethod
+    def setUpClass(cls):
+        """Run setup for Xenial to Bionic Series Upgrades."""
+        super(XenialBionicSeriesUpgradeUbuntu, cls).setUpClass()
+        cls.from_series = "xenial"
+        cls.to_series = "bionic"
+
+
+class BionicFocalSeriesUpgradeUbuntu(UbuntuLiteParallelSeriesUpgrade):
+    """OpenStack Bionic to FocalSeries Upgrade."""
+
+    @classmethod
+    def setUpClass(cls):
+        """Run setup for Xenial to Bionic Series Upgrades."""
+        super(BionicFocalSeriesUpgradeUbuntu, cls).setUpClass()
+        cls.from_series = "bionic"
+        cls.to_series = "focal"
+
+
 if __name__ == "__main__":
     from_series = os.environ.get("FROM_SERIES")
     if from_series == "trusty":

From fa4587f3668aa1adcd7e4dfadff01c8d5768f848 Mon Sep 17 00:00:00 2001
From: Chris MacNaughton <chris.macnaughton@canonical.com>
Date: Wed, 15 Apr 2020 16:34:21 +0200
Subject: [PATCH 17/19] Ensure that origin is set only after the first machine
 is rebooting

---
 ..._zaza_utilities_parallel_series_upgrade.py | 71 +++++++++++++++----
 .../utilities/parallel_series_upgrade.py      | 18 +++--
 2 files changed, 69 insertions(+), 20 deletions(-)

diff --git a/unit_tests/utilities/test_zaza_utilities_parallel_series_upgrade.py b/unit_tests/utilities/test_zaza_utilities_parallel_series_upgrade.py
index 36638b1..616c590 100644
--- a/unit_tests/utilities/test_zaza_utilities_parallel_series_upgrade.py
+++ b/unit_tests/utilities/test_zaza_utilities_parallel_series_upgrade.py
@@ -200,7 +200,6 @@ class TestParallelSeriesUpgrade(AioTestCase):
         mock_get_class.assert_called_once_with('my.thing')
         called.assert_called_once_with('1')
 
-    @mock.patch.object(upgrade_utils.os_utils, 'async_set_origin')
     @mock.patch.object(upgrade_utils, 'run_post_application_upgrade_functions')
     @mock.patch.object(
         upgrade_utils.series_upgrade_utils, 'async_prepare_series_upgrade')
@@ -214,7 +213,6 @@ class TestParallelSeriesUpgrade(AioTestCase):
         mock_async_set_series,
         mock_async_prepare_series_upgrade,
         mock_post_application_upgrade_functions,
-        mock_async_set_origin,
     ):
         self.juju_status.return_value.applications.__getitem__.return_value = \
             FAKE_STATUS_MONGO
@@ -240,24 +238,28 @@ class TestParallelSeriesUpgrade(AioTestCase):
         mock_series_upgrade_machine.assert_has_calls([
             mock.call(
                 '1',
+                origin=None,
+                application='mongodb',
                 files=None,
                 workaround_script=None,
                 post_upgrade_functions=[]),
             mock.call(
                 '2',
+                origin=None,
+                application='mongodb',
                 files=None,
                 workaround_script=None,
                 post_upgrade_functions=[]),
             mock.call(
                 '0',
+                origin=None,
+                application='mongodb',
                 files=None,
                 workaround_script=None,
                 post_upgrade_functions=[]),
         ])
-        mock_async_set_origin.assert_not_called()
         mock_post_application_upgrade_functions.assert_called_once_with([])
 
-    @mock.patch.object(upgrade_utils.os_utils, 'async_set_origin')
     @mock.patch.object(upgrade_utils, 'run_post_application_upgrade_functions')
     @mock.patch.object(
         upgrade_utils.series_upgrade_utils, 'async_prepare_series_upgrade')
@@ -271,7 +273,6 @@ class TestParallelSeriesUpgrade(AioTestCase):
         mock_async_set_series,
         mock_async_prepare_series_upgrade,
         mock_post_application_upgrade_functions,
-        mock_async_set_origin,
     ):
         self.juju_status.return_value.applications.__getitem__.return_value = \
             FAKE_STATUS_MONGO
@@ -294,24 +295,28 @@ class TestParallelSeriesUpgrade(AioTestCase):
         mock_series_upgrade_machine.assert_has_calls([
             mock.call(
                 '1',
+                origin=None,
+                application='mongodb',
                 files=None,
                 workaround_script=None,
                 post_upgrade_functions=[]),
             mock.call(
                 '2',
+                origin=None,
+                application='mongodb',
                 files=None,
                 workaround_script=None,
                 post_upgrade_functions=[]),
             mock.call(
                 '0',
+                origin=None,
+                application='mongodb',
                 files=None,
                 workaround_script=None,
                 post_upgrade_functions=[]),
         ])
-        mock_async_set_origin.assert_not_called()
         mock_post_application_upgrade_functions.assert_called_once_with([])
 
-    @mock.patch.object(upgrade_utils.os_utils, 'async_set_origin')
     @mock.patch.object(upgrade_utils, 'run_post_application_upgrade_functions')
     @mock.patch.object(
         upgrade_utils.series_upgrade_utils, 'async_prepare_series_upgrade')
@@ -325,7 +330,6 @@ class TestParallelSeriesUpgrade(AioTestCase):
         mock_async_set_series,
         mock_async_prepare_series_upgrade,
         mock_post_application_upgrade_functions,
-        mock_async_set_origin,
     ):
         await upgrade_utils.parallel_series_upgrade(
             'app',
@@ -335,34 +339,39 @@ class TestParallelSeriesUpgrade(AioTestCase):
         mock_async_set_series.assert_called_once_with(
             'app', to_series='xenial')
         self.juju_status.assert_called()
+        # The below is using `any_order=True` because the ordering is
+        # undetermined and differs between python versions
         mock_async_prepare_series_upgrade.assert_has_calls([
             mock.call('1', to_series='xenial'),
             mock.call('2', to_series='xenial'),
             mock.call('0', to_series='xenial'),
-        ])
+        ], any_order=True)
         mock_maybe_pause_things.assert_called()
         mock_series_upgrade_machine.assert_has_calls([
             mock.call(
                 '1',
+                origin='openstack-origin',
+                application='app',
                 files=None,
                 workaround_script=None,
                 post_upgrade_functions=None),
             mock.call(
                 '2',
+                origin='openstack-origin',
+                application='app',
                 files=None,
                 workaround_script=None,
                 post_upgrade_functions=None),
             mock.call(
                 '0',
+                origin='openstack-origin',
+                application='app',
                 files=None,
                 workaround_script=None,
                 post_upgrade_functions=None),
         ])
-        mock_async_set_origin.assert_called_once_with(
-            'app', 'openstack-origin')
         mock_post_application_upgrade_functions.assert_called_once_with(None)
 
-    @mock.patch.object(upgrade_utils.os_utils, 'async_set_origin')
     @mock.patch.object(upgrade_utils, 'run_post_application_upgrade_functions')
     @mock.patch.object(
         upgrade_utils.series_upgrade_utils, 'async_prepare_series_upgrade')
@@ -376,7 +385,6 @@ class TestParallelSeriesUpgrade(AioTestCase):
         mock_async_set_series,
         mock_async_prepare_series_upgrade,
         mock_post_application_upgrade_functions,
-        mock_async_set_origin,
     ):
         await upgrade_utils.serial_series_upgrade(
             'app',
@@ -395,22 +403,26 @@ class TestParallelSeriesUpgrade(AioTestCase):
         mock_series_upgrade_machine.assert_has_calls([
             mock.call(
                 '0',
+                origin='openstack-origin',
+                application='app',
                 files=None,
                 workaround_script=None,
                 post_upgrade_functions=None),
             mock.call(
                 '1',
+                origin='openstack-origin',
+                application='app',
                 files=None,
                 workaround_script=None,
                 post_upgrade_functions=None),
             mock.call(
                 '2',
+                origin='openstack-origin',
+                application='app',
                 files=None,
                 workaround_script=None,
                 post_upgrade_functions=None),
         ])
-        mock_async_set_origin.assert_called_once_with(
-            'app', 'openstack-origin')
         mock_post_application_upgrade_functions.assert_called_once_with(None)
 
     @mock.patch.object(
@@ -436,6 +448,35 @@ class TestParallelSeriesUpgrade(AioTestCase):
         mock_reboot.assert_called_once_with('1')
         mock_async_complete_series_upgrade.assert_called_once_with('1')
 
+    @mock.patch.object(upgrade_utils.os_utils, 'async_set_origin')
+    @mock.patch.object(
+        upgrade_utils.series_upgrade_utils, 'async_complete_series_upgrade')
+    @mock.patch.object(upgrade_utils, 'reboot')
+    @mock.patch.object(upgrade_utils, 'async_do_release_upgrade')
+    @mock.patch.object(upgrade_utils, 'async_dist_upgrade')
+    async def test_series_upgrade_machine_with_source(
+        self,
+        mock_async_dist_upgrade,
+        mock_async_do_release_upgrade,
+        mock_reboot,
+        mock_async_complete_series_upgrade,
+        mock_async_set_origin
+    ):
+        await upgrade_utils.series_upgrade_machine(
+            '1',
+            origin='openstack-origin',
+            application='app',
+            post_upgrade_functions=None,
+            pre_upgrade_functions=None,
+            files=None,
+            workaround_script=None)
+        mock_async_dist_upgrade.assert_called_once_with('1')
+        mock_async_do_release_upgrade.assert_called_once_with('1')
+        mock_reboot.assert_called_once_with('1')
+        mock_async_complete_series_upgrade.assert_called_once_with('1')
+        mock_async_set_origin.assert_called_once_with(
+            'app', 'openstack-origin')
+
     async def test_maybe_pause_things_primary(self):
         await upgrade_utils.maybe_pause_things(
             FAKE_STATUS,
diff --git a/zaza/openstack/utilities/parallel_series_upgrade.py b/zaza/openstack/utilities/parallel_series_upgrade.py
index dfbd36e..980f556 100755
--- a/zaza/openstack/utilities/parallel_series_upgrade.py
+++ b/zaza/openstack/utilities/parallel_series_upgrade.py
@@ -207,8 +207,6 @@ async def parallel_series_upgrade(
         pause_non_leader_primary)
     await series_upgrade_utils.async_set_series(
         application, to_series=to_series)
-    if origin:
-        await os_utils.async_set_origin(application, origin)
     app_idle = [
         wait_for_unit_idle(unit) for unit in status["units"]
     ]
@@ -223,6 +221,8 @@ async def parallel_series_upgrade(
     upgrade_group = [
         series_upgrade_machine(
             machine,
+            origin=origin,
+            application=application,
             files=files, workaround_script=workaround_script,
             post_upgrade_functions=post_upgrade_functions)
         for machine in machines
@@ -308,8 +308,6 @@ async def serial_series_upgrade(
         pause_non_leader_primary)
     await series_upgrade_utils.async_set_series(
         application, to_series=to_series)
-    if origin:
-        await os_utils.async_set_origin(application, origin)
     if not follower_first and leader_machine not in completed_machines:
         await wait_for_unit_idle(leader)
         await prepare_series_upgrade(leader_machine, to_series=to_series)
@@ -317,6 +315,8 @@ async def serial_series_upgrade(
                      .format(application, leader_machine))
         await series_upgrade_machine(
             leader_machine,
+            origin=origin,
+            application=application,
             files=files, workaround_script=workaround_script,
             post_upgrade_functions=post_upgrade_functions)
         completed_machines.append(leader_machine)
@@ -332,6 +332,8 @@ async def serial_series_upgrade(
                      .format(application, machine))
         await series_upgrade_machine(
             machine,
+            origin=origin,
+            application=application,
             files=files, workaround_script=workaround_script,
             post_upgrade_functions=post_upgrade_functions)
         completed_machines.append(machine)
@@ -343,6 +345,8 @@ async def serial_series_upgrade(
                      .format(application, leader_machine))
         await series_upgrade_machine(
             leader_machine,
+            origin=origin,
+            application=application,
             files=files, workaround_script=workaround_script,
             post_upgrade_functions=post_upgrade_functions)
         completed_machines.append(leader_machine)
@@ -352,6 +356,8 @@ async def serial_series_upgrade(
 
 async def series_upgrade_machine(
         machine,
+        origin=None,
+        application=None,
         post_upgrade_functions=None,
         pre_upgrade_functions=None,
         files=None,
@@ -379,6 +385,8 @@ async def series_upgrade_machine(
     await async_dist_upgrade(machine)
     await async_do_release_upgrade(machine)
     await reboot(machine)
+    if origin:
+        await os_utils.async_set_origin(application, origin)
     await series_upgrade_utils.async_complete_series_upgrade(machine)
     await run_post_upgrade_functions(post_upgrade_functions)
 
@@ -462,8 +470,8 @@ async def maybe_pause_things(
             logging.info("Pausing {}".format(unit))
             leader_pauses.append(
                 model.async_run_action(unit, "pause", action_params={}))
-    await asyncio.gather(*subordinate_pauses)
     await asyncio.gather(*leader_pauses)
+    await asyncio.gather(*subordinate_pauses)
 
 
 def get_leader_and_non_leaders(status):

From 2bc2234cad71b31253d8d17e59a1abb3be590f56 Mon Sep 17 00:00:00 2001
From: Chris MacNaughton <chris.macnaughton@canonical.com>
Date: Thu, 16 Apr 2020 11:53:35 +0200
Subject: [PATCH 18/19] Update to migrate bits out to Zaza

---
 ..._zaza_utilities_parallel_series_upgrade.py | 33 ++------
 .../utilities/parallel_series_upgrade.py      | 84 ++-----------------
 2 files changed, 18 insertions(+), 99 deletions(-)

diff --git a/unit_tests/utilities/test_zaza_utilities_parallel_series_upgrade.py b/unit_tests/utilities/test_zaza_utilities_parallel_series_upgrade.py
index 616c590..aefa00b 100644
--- a/unit_tests/utilities/test_zaza_utilities_parallel_series_upgrade.py
+++ b/unit_tests/utilities/test_zaza_utilities_parallel_series_upgrade.py
@@ -179,6 +179,9 @@ class TestParallelSeriesUpgrade(AioTestCase):
 
         self.async_block_until = mock.AsyncMock()
         self.model.async_block_until = self.async_block_until
+        self.model.async_wait_for_unit_idle = mock.AsyncMock()
+        self.async_run_on_machine = mock.AsyncMock()
+        self.model.async_run_on_machine = self.async_run_on_machine
 
     @mock.patch.object(upgrade_utils.cl_utils, 'get_class')
     async def test_run_post_application_upgrade_functions(
@@ -520,13 +523,12 @@ class TestParallelSeriesUpgrade(AioTestCase):
             pause_non_leader_primary=False)
         self.async_run_action.assert_not_called()
 
-    @mock.patch.object(upgrade_utils, 'run_on_machine')
-    async def test_async_do_release_upgrade(self, mock_run_on_machine):
+    async def test_async_do_release_upgrade(self):
         await upgrade_utils.async_do_release_upgrade('1')
         do_release_upgrade_cmd = (
             'yes | sudo DEBIAN_FRONTEND=noninteractive '
             'do-release-upgrade -f DistUpgradeViewNonInteractive')
-        mock_run_on_machine.assert_called_once_with(
+        self.async_run_on_machine.assert_called_once_with(
             '1', do_release_upgrade_cmd, timeout='120m'
         )
 
@@ -538,37 +540,20 @@ class TestParallelSeriesUpgrade(AioTestCase):
             '1', to_series='xenial'
         )
 
-    @mock.patch.object(upgrade_utils, 'run_on_machine')
-    async def test_reboot(self, mock_run_on_machine):
+    async def test_reboot(self):
         await upgrade_utils.reboot('1')
-        mock_run_on_machine.assert_called_once_with(
+        self.async_run_on_machine.assert_called_once_with(
             '1', 'sudo init 6 & exit'
         )
 
-    async def test_run_on_machine(self):
-        await upgrade_utils.run_on_machine('1', 'test')
-        self.check_call.assert_called_once_with(
-            ['juju', 'run', '--machine=1', 'test'])
-
-    async def test_run_on_machine_with_timeout(self):
-        await upgrade_utils.run_on_machine('1', 'test', timeout='20m')
-        self.check_call.assert_called_once_with(
-            ['juju', 'run', '--machine=1', '--timeout=20m', 'test'])
-
-    async def test_run_on_machine_with_model(self):
-        await upgrade_utils.run_on_machine('1', 'test', model_name='test')
-        self.check_call.assert_called_once_with(
-            ['juju', 'run', '--machine=1', '--model=test', 'test'])
-
-    @mock.patch.object(upgrade_utils, 'run_on_machine')
-    async def test_async_dist_upgrade(self, mock_run_on_machine):
+    async def test_async_dist_upgrade(self):
         await upgrade_utils.async_dist_upgrade('1')
         apt_update_command = (
             """yes | sudo DEBIAN_FRONTEND=noninteractive """
             """apt-get --assume-yes """
             """-o "Dpkg::Options::=--force-confdef" """
             """-o "Dpkg::Options::=--force-confold" dist-upgrade""")
-        mock_run_on_machine.assert_has_calls([
+        self.async_run_on_machine.assert_has_calls([
             mock.call('1', 'sudo apt-get update'),
             mock.call('1', apt_update_command),
         ])
diff --git a/zaza/openstack/utilities/parallel_series_upgrade.py b/zaza/openstack/utilities/parallel_series_upgrade.py
index 980f556..eea2edb 100755
--- a/zaza/openstack/utilities/parallel_series_upgrade.py
+++ b/zaza/openstack/utilities/parallel_series_upgrade.py
@@ -18,7 +18,6 @@
 
 import asyncio
 
-import concurrent
 import collections
 import copy
 import logging
@@ -208,7 +207,7 @@ async def parallel_series_upgrade(
     await series_upgrade_utils.async_set_series(
         application, to_series=to_series)
     app_idle = [
-        wait_for_unit_idle(unit) for unit in status["units"]
+        model.async_wait_for_unit_idle(unit) for unit in status["units"]
     ]
     await asyncio.gather(*app_idle)
     prepare_group = [
@@ -309,7 +308,7 @@ async def serial_series_upgrade(
     await series_upgrade_utils.async_set_series(
         application, to_series=to_series)
     if not follower_first and leader_machine not in completed_machines:
-        await wait_for_unit_idle(leader)
+        await model.async_wait_for_unit_idle(leader)
         await prepare_series_upgrade(leader_machine, to_series=to_series)
         logging.info("About to upgrade leader of {}: {}"
                      .format(application, leader_machine))
@@ -326,7 +325,7 @@ async def serial_series_upgrade(
         machine = unit['machine']
         if machine in completed_machines:
             continue
-        await wait_for_unit_idle(unit_name)
+        await model.async_wait_for_unit_idle(unit_name)
         await prepare_series_upgrade(machine, to_series=to_series)
         logging.info("About to upgrade follower of {}: {}"
                      .format(application, machine))
@@ -339,7 +338,7 @@ async def serial_series_upgrade(
         completed_machines.append(machine)
 
     if follower_first and leader_machine not in completed_machines:
-        await wait_for_unit_idle(leader)
+        await model.async_wait_for_unit_idle(leader)
         await prepare_series_upgrade(leader_machine, to_series=to_series)
         logging.info("About to upgrade leader of {}: {}"
                      .format(application, leader_machine))
@@ -530,37 +529,13 @@ async def reboot(machine):
     :rtype: None
     """
     try:
-        await run_on_machine(machine, 'sudo init 6 & exit')
+        await model.async_run_on_machine(machine, 'sudo init 6 & exit')
         # await run_on_machine(unit, "sudo reboot && exit")
     except subprocess.CalledProcessError as e:
         logging.warn("Error doing reboot: {}".format(e))
         pass
 
 
-async def run_on_machine(machine, command, model_name=None, timeout=None):
-    """Juju run on unit.
-
-    :param model_name: Name of model unit is in
-    :type model_name: str
-    :param unit_name: Name of unit to match
-    :type unit: str
-    :param command: Command to execute
-    :type command: str
-    :param timeout: How long in seconds to wait for command to complete
-    :type timeout: int
-    :returns: action.data['results'] {'Code': '', 'Stderr': '', 'Stdout': ''}
-    :rtype: dict
-    """
-    cmd = ['juju', 'run', '--machine={}'.format(machine)]
-    if model_name:
-        cmd.append('--model={}'.format(model_name))
-    if timeout:
-        cmd.append('--timeout={}'.format(timeout))
-    cmd.append(command)
-    logging.info("About to call '{}'".format(cmd))
-    await os_utils.check_call(cmd)
-
-
 async def async_dist_upgrade(machine):
     """Run dist-upgrade on unit after update package db.
 
@@ -571,14 +546,14 @@ async def async_dist_upgrade(machine):
     """
     logging.info('Updating package db ' + machine)
     update_cmd = 'sudo apt-get update'
-    await run_on_machine(machine, update_cmd)
+    await model.async_run_on_machine(machine, update_cmd)
 
     logging.info('Updating existing packages ' + machine)
     dist_upgrade_cmd = (
         """yes | sudo DEBIAN_FRONTEND=noninteractive apt-get --assume-yes """
         """-o "Dpkg::Options::=--force-confdef" """
         """-o "Dpkg::Options::=--force-confold" dist-upgrade""")
-    await run_on_machine(machine, dist_upgrade_cmd)
+    await model.async_run_on_machine(machine, dist_upgrade_cmd)
 
 
 async def async_do_release_upgrade(machine):
@@ -594,46 +569,5 @@ async def async_do_release_upgrade(machine):
         'yes | sudo DEBIAN_FRONTEND=noninteractive '
         'do-release-upgrade -f DistUpgradeViewNonInteractive')
 
-    await run_on_machine(machine, do_release_upgrade_cmd, timeout="120m")
-
-
-# TODO: Move these functions into zaza.model
-async def wait_for_unit_idle(unit_name, timeout=600):
-    """Wait until the unit's agent is idle.
-
-    :param unit_name: The unit name of the application, ex: mysql/0
-    :type unit_name: str
-    :param timeout: How long to wait before timing out
-    :type timeout: int
-    :returns: None
-    :rtype: None
-    """
-    app = unit_name.split('/')[0]
-    try:
-        await model.async_block_until(
-            _unit_idle(app, unit_name),
-            timeout=timeout)
-    except concurrent.futures._base.TimeoutError:
-        raise model.ModelTimeout("Zaza has timed out waiting on {} to "
-                                 "reach idle state.".format(unit_name))
-
-
-def _unit_idle(app, unit_name):
-    async def f():
-        x = await get_agent_status(app, unit_name)
-        return x == "idle"
-    return f
-
-
-async def get_agent_status(app, unit_name):
-    """Get the current status of the specified unit.
-
-    :param app: The name of the Juju application, ex: mysql
-    :type app: str
-    :param unit_name: The unit name of the application, ex: mysql/0
-    :type unit_name: str
-    :returns: The agent status, either active / idle, returned by Juju
-    :rtype: str
-    """
-    return (await model.async_get_status()). \
-        applications[app]['units'][unit_name]['agent-status']['status']
+    await model.async_run_on_machine(
+        machine, do_release_upgrade_cmd, timeout="120m")

From de028d7eb2b9fe1b35b43a8796c284fdf91be537 Mon Sep 17 00:00:00 2001
From: Chris MacNaughton <chris.macnaughton@canonical.com>
Date: Thu, 16 Apr 2020 11:58:54 +0200
Subject: [PATCH 19/19] fix copy-pasta error

---
 zaza/openstack/utilities/parallel_series_upgrade.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/zaza/openstack/utilities/parallel_series_upgrade.py b/zaza/openstack/utilities/parallel_series_upgrade.py
index eea2edb..29e276d 100755
--- a/zaza/openstack/utilities/parallel_series_upgrade.py
+++ b/zaza/openstack/utilities/parallel_series_upgrade.py
@@ -247,7 +247,7 @@ async def serial_series_upgrade(
     files=None,
     workaround_script=None
 ):
-    """Perform series upgrade on an application in parallel.
+    """Perform series upgrade on an application in serial.
 
     :param unit_name: Unit Name
     :type unit_name: str