From f3ebe916f282c080d50501e3edd895a40d3d15b8 Mon Sep 17 00:00:00 2001
From: Liam Young <liam.young@canonical.com>
Date: Wed, 11 Dec 2019 12:12:25 +0000
Subject: [PATCH] Add get_subordinate_units

Add zaza.openstack.utilities.juju.get_subordinate_units to get a list
of all subordinate units associated with units in unit_list.
Subordinate can be filtered by using 'charm_name' which will only
return subordinate units which have 'charm_name' in the name of the
charm.
---
 .../utilities/test_zaza_utilities_juju.py     | 31 +++++++++++++
 zaza/openstack/utilities/juju.py              | 46 +++++++++++++++++++
 2 files changed, 77 insertions(+)

diff --git a/unit_tests/utilities/test_zaza_utilities_juju.py b/unit_tests/utilities/test_zaza_utilities_juju.py
index 11bbcbe..4255f9e 100644
--- a/unit_tests/utilities/test_zaza_utilities_juju.py
+++ b/unit_tests/utilities/test_zaza_utilities_juju.py
@@ -279,3 +279,34 @@ class TestJujuUtils(ut_utils.BaseTestCase):
             model_name=None
         )
         self.assertEqual(expected, actual)
+
+    def test_get_subordinate_units(self):
+        juju_status = mock.MagicMock()
+        juju_status.applications = {
+            'nova-compute': {
+                'units': {
+                    'nova-compute/0': {
+                        'subordinates': {
+                            'neutron-openvswitch/2': {
+                                'charm': 'cs:neutron-openvswitch-22'}}}}},
+            'cinder': {
+                'units': {
+                    'cinder/1': {
+                        'subordinates': {
+                            'cinder-hacluster/0': {
+                                'charm': 'cs:hacluster-42'},
+                            'cinder-ceph/3': {
+                                'charm': 'cs:cinder-ceph-2'}}}}},
+        }
+        self.assertEqual(
+            sorted(juju_utils.get_subordinate_units(
+                ['nova-compute/0', 'cinder/1'],
+                status=juju_status)),
+            sorted(['neutron-openvswitch/2', 'cinder-hacluster/0',
+                    'cinder-ceph/3']))
+        self.assertEqual(
+            juju_utils.get_subordinate_units(
+                ['nova-compute/0', 'cinder/1'],
+                charm_name='ceph',
+                status=juju_status),
+            ['cinder-ceph/3'])
diff --git a/zaza/openstack/utilities/juju.py b/zaza/openstack/utilities/juju.py
index fd7abc1..99f1e10 100644
--- a/zaza/openstack/utilities/juju.py
+++ b/zaza/openstack/utilities/juju.py
@@ -311,3 +311,49 @@ def leader_get(application, key='', model_name=None):
         return yaml.safe_load(result.get('Stdout'))
     else:
         raise model.CommandRunFailed(cmd, result)
+
+
+def get_subordinate_units(unit_list, charm_name=None, status=None,
+                          model_name=None):
+    """Get a list of all subordinate units associated with units in unit_list.
+
+    Get a list of all subordinate units associated with units in unit_list.
+    Subordinate can be filtered by using 'charm_name' which will only return
+    subordinate units which have 'charm_name' in the name of the charm e.g.
+
+        get_subordinate_units(
+            ['cinder/1']) would return ['cinder-hacluster/1',
+                                        'cinder-ceph/2'])
+    where as
+
+        get_subordinate_units(
+            ['cinder/1'], charm_name='hac') would return ['cinder-hacluster/1']
+
+    NOTE: The charm_name match is against the name of the charm not the
+          application name.
+
+    :param charm_name: List of unit names
+    :type unit_list: []
+    :param charm_name: charm_name to match against, can be a sub-string.
+    :type charm_name: str
+    :param status: Juju status to query against,
+    :type status: juju.client._definitions.FullStatus
+    :param model_name: Name of model to query.
+    :type model_name: str
+    :returns: List of matching unit names.
+    :rtype: []
+    """
+    if not status:
+        status = model.get_status(model_name=model_name)
+    sub_units = []
+    for unit_name in unit_list:
+        app_name = unit_name.split('/')[0]
+        subs = status.applications[app_name]['units'][unit_name].get(
+            'subordinates') or {}
+        if charm_name:
+            for unit_name, unit_data in subs.items():
+                if charm_name in unit_data['charm']:
+                    sub_units.append(unit_name)
+        else:
+            sub_units.extend([n for n in subs.keys()])
+    return sub_units