From 5d7d622b026f6f4c6994fc5fbb1789d7827bf641 Mon Sep 17 00:00:00 2001
From: David Garcia <david.garcia@canonical.com>
Date: Wed, 10 Mar 2021 11:50:38 +0100
Subject: [PATCH] Fix bug 1448: minor fix in JujuModelWatcher.wait_for_model

Add an `application_ready` function in the Juju Watcher to check if all the
applications of the model are ready.  This function checks that not only
the application is ready, but also all the units of the application.

Additionally, I have added a stability check. After considering a model
as ready, we will wait for 10 seconds, and check again.

Change-Id: Icd8d88613416b75cfd241e7044dbeffa6cbc6a8e
Signed-off-by: David Garcia <david.garcia@canonical.com>
(cherry picked from commit c344117335d4ffbd06ef90ae8dce9cb3910165fb)
---
 n2vc/juju_watcher.py | 68 ++++++++++++++++++++++++++++++++------------
 1 file changed, 50 insertions(+), 18 deletions(-)

diff --git a/n2vc/juju_watcher.py b/n2vc/juju_watcher.py
index e122786..04ad10f 100644
--- a/n2vc/juju_watcher.py
+++ b/n2vc/juju_watcher.py
@@ -35,6 +35,13 @@ def status(application: Application) -> str:
 
 
 def entity_ready(entity: ModelEntity) -> bool:
+    """
+    Check if the entity is ready
+
+    :param: entity: Model entity. It can be a machine, action, or application.
+
+    :returns: boolean saying if the entity is ready or not
+    """
     entity_type = entity.entity_type
     if entity_type == "machine":
         return entity.agent_status in ["started"]
@@ -42,11 +49,27 @@ def entity_ready(entity: ModelEntity) -> bool:
         return entity.status in ["completed", "failed", "cancelled"]
     elif entity_type == "application":
         # Workaround for bug: https://github.com/juju/python-libjuju/issues/441
-        return status(entity) in ["active", "blocked"]
+        return entity.status in ["active", "blocked"]
     else:
         raise EntityInvalidException("Unknown entity type: {}".format(entity_type))
 
 
+def application_ready(application: Application) -> bool:
+    """
+    Check if an application has a leader
+
+    :param: application: Application entity.
+
+    :returns: boolean saying if the application has a unit that is a leader.
+    """
+    ready_status_list = ["active", "blocked"]
+    application_ready = application.status in ready_status_list
+    units_ready = all(
+        unit.workload_status in ready_status_list for unit in application.units
+    )
+    return application_ready and units_ready
+
+
 class JujuModelWatcher:
     @staticmethod
     async def wait_for_model(model: Model, timeout: float = 3600):
@@ -63,24 +86,32 @@ class JujuModelWatcher:
             timeout = 3600.0
 
         # Coroutine to wait until the entity reaches the final state
-        wait_for_entity = asyncio.ensure_future(
-            asyncio.wait_for(
-                model.block_until(
-                    lambda: all(
-                        entity_ready(entity) for entity in model.applications.values()
-                    )
-                ),
-                timeout=timeout,
+        async def wait_until_model_ready():
+            wait_for_entity = asyncio.ensure_future(
+                asyncio.wait_for(
+                    model.block_until(
+                        lambda: all(
+                            application_ready(application)
+                            for application in model.applications.values()
+                        ),
+                    ),
+                    timeout=timeout,
+                )
             )
-        )
 
-        tasks = [wait_for_entity]
-        try:
-            await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
-        finally:
-            # Cancel tasks
-            for task in tasks:
-                task.cancel()
+            tasks = [wait_for_entity]
+            try:
+                await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
+            finally:
+                # Cancel tasks
+                for task in tasks:
+                    task.cancel()
+
+        await wait_until_model_ready()
+        # Check model is still ready after 10 seconds
+
+        await asyncio.sleep(10)
+        await wait_until_model_ready()
 
     @staticmethod
     async def wait_for(
@@ -116,7 +147,8 @@ class JujuModelWatcher:
         # Coroutine to wait until the entity reaches the final state
         wait_for_entity = asyncio.ensure_future(
             asyncio.wait_for(
-                model.block_until(lambda: entity_ready(entity)), timeout=total_timeout,
+                model.block_until(lambda: entity_ready(entity)),
+                timeout=total_timeout,
             )
         )
 
-- 
2.17.1