Fix bug 1298

[osm/N2VC.git] / n2vc / juju_watcher.py

diff --git a/n2vc/juju_watcher.py b/n2vc/juju_watcher.py
index 842e990..e122786 100644 (file)
--- a/n2vc/juju_watcher.py
+++ b/n2vc/juju_watcher.py
@@ -48,9 +48,43 @@ def entity_ready(entity: ModelEntity) -> bool:
 
 
 class JujuModelWatcher:
+    @staticmethod
+    async def wait_for_model(model: Model, timeout: float = 3600):
+        """
+        Wait for all entities in model to reach its final state.
+
+        :param: model:              Model to observe
+        :param: timeout:            Timeout for the model applications to be active
+
+        :raises: asyncio.TimeoutError when timeout reaches
+        """
+
+        if timeout is None:
+            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,
+            )
+        )
+
+        tasks = [wait_for_entity]
+        try:
+            await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
+        finally:
+            # Cancel tasks
+            for task in tasks:
+                task.cancel()
+
     @staticmethod
     async def wait_for(
-        model,
+        model: Model,
         entity: ModelEntity,
         progress_timeout: float = 3600,
         total_timeout: float = 3600,
@@ -103,8 +137,6 @@ class JujuModelWatcher:
             # Execute tasks, and stop when the first is finished
             # The watcher task won't never finish (unless it timeouts)
             await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
-        except Exception as e:
-            raise e
         finally:
             # Cancel tasks
             for task in tasks: