Wait for cloud-init to finish before provisioning

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

diff --git a/n2vc/libjuju.py b/n2vc/libjuju.py
index a6fd8fe..ce1c9df 100644 (file)
--- a/n2vc/libjuju.py
+++ b/n2vc/libjuju.py
@@ -273,7 +273,9 @@ class Libjuju:
         """
         return await controller.get_model(model_name)
 
-    async def model_exists(self, model_name: str, controller: Controller = None) -> bool:
+    async def model_exists(
+        self, model_name: str, controller: Controller = None
+    ) -> bool:
         """
         Check if model exists
 
@@ -661,6 +663,83 @@ class Libjuju:
 
         return application
 
+    async def scale_application(
+        self,
+        model_name: str,
+        application_name: str,
+        scale: int = 1,
+        total_timeout: float = None,
+    ):
+        """
+        Scale application (K8s)
+
+        :param: model_name:         Model name
+        :param: application_name:   Application name
+        :param: scale:              Scale to which to set this application
+        :param: total_timeout:      Timeout for the entity to be active
+        """
+
+        model = None
+        controller = await self.get_controller()
+        try:
+            model = await self.get_model(controller, model_name)
+
+            self.log.debug(
+                "Scaling application {} in model {}".format(
+                    application_name, model_name
+                )
+            )
+            application = self._get_application(model, application_name)
+            if application is None:
+                raise JujuApplicationNotFound("Cannot scale application")
+            await application.scale(scale=scale)
+            # Wait until application is scaled in model
+            self.log.debug(
+                "Waiting for application {} to be scaled in model {}...".format(
+                    application_name, model_name
+                )
+            )
+            if total_timeout is None:
+                total_timeout = 1800
+            end = time.time() + total_timeout
+            while time.time() < end:
+                application_scale = self._get_application_count(model, application_name)
+                # Before calling wait_for_model function,
+                # wait until application unit count and scale count are equal.
+                # Because there is a delay before scaling triggers in Juju model.
+                if application_scale == scale:
+                    await JujuModelWatcher.wait_for_model(
+                        model=model, timeout=total_timeout
+                    )
+                    self.log.debug(
+                        "Application {} is scaled in model {}".format(
+                            application_name, model_name
+                        )
+                    )
+                    return
+                await asyncio.sleep(5)
+            raise Exception(
+                "Timeout waiting for application {} in model {} to be scaled".format(
+                    application_name, model_name
+                )
+            )
+        finally:
+            if model:
+                await self.disconnect_model(model)
+            await self.disconnect_controller(controller)
+
+    def _get_application_count(self, model: Model, application_name: str) -> int:
+        """Get number of units of the application
+
+        :param: model:              Model object
+        :param: application_name:   Application name
+
+        :return: int (or None if application doesn't exist)
+        """
+        application = self._get_application(model, application_name)
+        if application is not None:
+            return len(application.units)
+
     def _get_application(self, model: Model, application_name: str) -> Application:
         """Get application