From fa31b98ef61fd38c16151676f11299ed8428e2e6 Mon Sep 17 00:00:00 2001
From: Lili Deng <lildeng@microsoft.com>
Date: Thu, 19 Sep 2024 19:58:29 +0800
Subject: [PATCH] extension: add wait_for_all_extensions_to_finish before
 installing extension

---
 lisa/sut_orchestrator/azure/features.py | 46 +++++++++++++++++++++++++
 1 file changed, 46 insertions(+)

diff --git a/lisa/sut_orchestrator/azure/features.py b/lisa/sut_orchestrator/azure/features.py
index efab4142b4..67e4235bb8 100644
--- a/lisa/sut_orchestrator/azure/features.py
+++ b/lisa/sut_orchestrator/azure/features.py
@@ -5,6 +5,7 @@
 import json
 import re
 import string
+import time
 from dataclasses import dataclass, field
 from functools import partial
 from pathlib import Path
@@ -2810,6 +2811,48 @@ def get(
         )
         return extension
 
+    def wait_for_all_extensions_to_finish(
+        self, wait_for_all_extensions_to_finish: bool = True
+    ) -> None:
+        if not wait_for_all_extensions_to_finish:
+            return
+
+        start_time = time.time()
+        max_wait_time = 120
+        check_interval = 10
+
+        while True:
+            extension_list = self.list_all()
+
+            if not extension_list.value:
+                self._log.debug("No extensions found for the VM. Exiting loop.")
+                break
+
+            all_extensions_completed = True
+
+            for extension in extension_list.value:
+                provisioning_state = extension.provisioning_state
+                self._log.debug(
+                    f"Provisioning state of {extension.name}: {provisioning_state}"
+                )
+
+                if provisioning_state not in ["Failed", "Succeeded"]:
+                    all_extensions_completed = False
+                    break
+
+            if all_extensions_completed:
+                self._log.debug("All extensions have completed successfully.")
+                break
+
+            elapsed_time = time.time() - start_time
+            if elapsed_time > max_wait_time:
+                self._log.error(
+                    f"Max wait time of {max_wait_time} seconds exceeded. Exiting loop."
+                )
+                break
+
+            time.sleep(check_interval)
+
     def create_or_update(
         self,
         type_: str,
@@ -2824,6 +2867,7 @@ def create_or_update(
         protected_settings: Any = None,
         suppress_failures: Optional[bool] = None,
         timeout: int = 60 * 25,
+        wait_for_all_extensions_to_finish: bool = True,
     ) -> Any:
         platform: AzurePlatform = self._platform  # type: ignore
         compute_client = get_compute_client(platform)
@@ -2851,6 +2895,8 @@ def create_or_update(
                 sub="***REDACTED***",
             )
 
+        self.wait_for_all_extensions_to_finish(wait_for_all_extensions_to_finish)
+
         self._log.debug(f"extension_parameters: {extension_parameters.as_dict()}")
 
         operation = compute_client.virtual_machine_extensions.begin_create_or_update(