From f4f25716a28bbc3effea5aff2b2779c9a2211ab8 Mon Sep 17 00:00:00 2001
From: Frode Nordahl <frode.nordahl@canonical.com>
Date: Mon, 20 Aug 2018 09:49:55 +0200
Subject: [PATCH] Add get_lead_unit_name function

---
 unit_tests/test_zaza_model.py | 10 ++++++++++
 zaza/model.py                 | 19 +++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/unit_tests/test_zaza_model.py b/unit_tests/test_zaza_model.py
index df5c068..e9be4fd 100644
--- a/unit_tests/test_zaza_model.py
+++ b/unit_tests/test_zaza_model.py
@@ -248,6 +248,16 @@ class TestModel(ut_utils.BaseTestCase):
             model.get_first_unit_name('model', 'app'),
             'app/2')
 
+    def test_get_lead_unit_name(self):
+        self.patch_object(model, 'get_juju_model', return_value='mname')
+        self.patch_object(model, 'get_units')
+        self.get_units.return_value = self.units
+        self.patch_object(model, 'Model')
+        self.Model.return_value = self.Model_mock
+        self.assertEqual(
+            model.get_lead_unit_name('app', 'model'),
+            'app/4')
+
     def test_get_unit_from_name(self):
         # Normal case
         self.patch_object(model, 'get_juju_model', return_value='mname')
diff --git a/zaza/model.py b/zaza/model.py
index 8656e9d..435613c 100644
--- a/zaza/model.py
+++ b/zaza/model.py
@@ -396,6 +396,25 @@ def get_first_unit_name(application_name, model_name=None):
     return get_units(application_name, model_name=model_name)[0].name
 
 
+async def async_get_lead_unit_name(application_name, model_name=None):
+    """Return name of lowest numbered unit of given application.
+
+    :param model_name: Name of model to query.
+    :type model_name: str
+    :param application_name: Name of application
+    :type application_name: str
+    :returns: Name of lowest numbered unit
+    :rtype: str
+    """
+    async with run_in_model(model_name) as model:
+        for unit in model.applications[application_name].units:
+            is_leader = await unit.is_leader_from_status()
+            if is_leader:
+                return unit.entity_id
+
+get_lead_unit_name = sync_wrapper(async_get_lead_unit_name)
+
+
 def get_app_ips(application_name, model_name=None):
     """Return public address of all units of an application.