Merge pull request #812 from vespa-engine/thomasht86/wait-for-prod-de…

…ployment

Thomasht86/wait for prod deployment

tests/integration/test_integration_vespa_cloud_token.py

@@ -1,7 +1,6 @@
 # Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 
 import os
-import time
 import asyncio
 import shutil
 import unittest
@@ -220,21 +219,11 @@ def setUp(self) -> None:
             source_url="https://github.com/vespa-engine/pyvespa",
         )
         # Wait for deployment to be ready
-        # Wait until buildstatus is succeeded
-        max_wait = 1800  # Could take up to 30 minutes
-        start = time.time()
-        success = False
-        while time.time() - start < max_wait:
-            build_status = self.vespa_cloud.check_production_build_status(
-                build_no=self.build_no
-            )
-            if build_status["status"] == "done":
-                if build_status["deployed"]:
-                    success = True
-                    break
-            time.sleep(5)
+        success = self.vespa_cloud.wait_for_prod_deployment(
+            build_no=self.build_no, max_wait=1800
+        )
         if not success:
-            raise ValueError("Deployment failed")
+            self.fail("Deployment failed")
         self.app = self.vespa_cloud.get_application(
             instance=self.instance_name, environment="prod"
         )

tests/integration/test_integration_vespa_cloud_vector_search.py

@@ -265,12 +265,12 @@ def setUp(self) -> None:
             )
         ]
 
-        # Validation override is needed to be able to use the certs generated by PyVespa, since they are overwritten each deployment
-        # tomorrow = datetime.now() + timedelta(days=1)
-        # formatted_date = tomorrow.strftime("%Y-%m-%d")
-        # self.app_package.validations = [
-        #     Validation(ValidationID("certificate-removal"), formatted_date)
-        # ]
+        # Validation override is needed to be able to be able to swith between triggering from local or CI.
+        tomorrow = datetime.now() + timedelta(days=1)
+        formatted_date = tomorrow.strftime("%Y-%m-%d")
+        self.app_package.validations = [
+            Validation(ValidationID("certificate-removal"), formatted_date)
+        ]
 
         self.vespa_cloud = VespaCloud(
             tenant="vespa-team",
@@ -285,21 +285,11 @@ def setUp(self) -> None:
             instance=self.instance_name,
             application_root=self.application_root,
         )
-        # Wait until buildstatus is succeeded
-        max_wait = 1800  # Could take up to 30 minutes
-        start = time.time()
-        success = False
-        while time.time() - start < max_wait:
-            build_status = self.vespa_cloud.check_production_build_status(
-                build_no=self.build_no
-            )
-            if build_status["status"] == "done":
-                if build_status["deployed"]:
-                    success = True
-                    break
-            time.sleep(5)
+        success = self.vespa_cloud.wait_for_prod_deployment(
+            build_no=self.build_no, max_wait=1800
+        )
         if not success:
-            raise ValueError("Deployment failed. Please check the Vespa Cloud console.")
+            self.fail("Deployment failed")
         self.app: Vespa = self.vespa_cloud.get_application(environment="prod")
         self.app.wait_for_application_up(max_wait=APP_INIT_TIMEOUT)
 

vespa/deployment.py

@@ -826,6 +826,39 @@ def check_production_build_status(self, build_no: Optional[int]) -> dict:
         )
         return status
 
+    def wait_for_prod_deployment(
+        self,
+        build_no: Optional[int] = None,
+        max_wait: int = 3600,
+        poll_interval: int = 5,
+    ) -> bool:
+        """
+        Wait for a production deployment to finish.
+        Useful for example in CI/CD pipelines to wait for a deployment to finish.
+
+        Example usage::
+
+            vespa_cloud = VespaCloud(...)
+            build_no = vespa_cloud.deploy_to_prod()
+            success = vespa_cloud.wait_for_prod_deployment(build_no, max_wait=3600, poll_interval=5)
+            print(success)
+            True
+
+        :param build_no: The build number to check.
+        :param max_wait: Maximum time to wait for the deployment in seconds. Default is 3600 (1 hour).
+        :param poll_interval: Polling interval in seconds. Default is 5 seconds.
+
+        :return: True if the deployment is done and converged. False if the deployment has failed.
+        :raises TimeoutError: If the deployment did not finish within max_wait seconds.
+        """
+        start_time = time.time()
+        while time.time() - start_time < max_wait:
+            status = self.check_production_build_status(build_no)
+            if status["status"] == "done":
+                return status["deployed"]
+            time.sleep(poll_interval)
+        raise TimeoutError(f"Deployment did not finish within {max_wait} seconds. ")
+
     def deploy_from_disk(
         self, instance: str, application_root: Path, max_wait: int = 300
     ) -> Vespa: