Use retried to wait for installation to finish

tests/integration/install/test_installation.py

@@ -1,7 +1,6 @@
 import dataclasses
 import json
 import logging
-import time
 from datetime import timedelta
 from typing import NoReturn
 
@@ -20,7 +19,6 @@
     AlreadyExists,
     InvalidParameterValue,
     NotFound,
-    ResourceDoesNotExist,
 )
 from databricks.sdk.retries import retried
 
@@ -394,24 +392,17 @@ def test_check_inventory_database_exists(ws, installation_ctx):
         installation_ctx.workspace_installer.configure()
 
 
-def test_compare_remote_local_install_versions(ws, installation_ctx):
+def test_compare_remote_local_install_versions(installation_ctx) -> None:
     installation_ctx.workspace_installation.run()
 
-    # max time to wait for the installation to finish
-    timeout_duration = 5
+    @retried(on=[NotFound], timeout=timedelta(minutes=2))
+    def wait_for_installation_to_finish():
+        installation_ctx.installation.load(WorkspaceConfig)
 
-    for _ in range(timeout_duration):
-        try:
-            installation_ctx.installation.load(WorkspaceConfig)
-            break
-        except ResourceDoesNotExist:
-            logger.info("Waiting for the installation to finish...")
-            time.sleep(1)
+    wait_for_installation_to_finish()
 
-    with pytest.raises(
-        RuntimeWarning,
-        match="UCX workspace remote and local install versions are same and no override is requested. Exiting...",
-    ):
+    error_message = "UCX workspace remote and local install versions are same and no override is requested. Exiting..."
+    with pytest.raises(RuntimeWarning, match=error_message):
         installation_ctx.workspace_installer.configure()
 
     installation_ctx.replace(