From 01d881066d0df91096b2d81eded2c2a78ab3666f Mon Sep 17 00:00:00 2001 From: pritishpai Date: Mon, 20 Jan 2025 19:07:06 +0530 Subject: [PATCH] Add timeout check --- tests/integration/install/test_installation.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/tests/integration/install/test_installation.py b/tests/integration/install/test_installation.py index a0b9b6f55b..68eceefb3a 100644 --- a/tests/integration/install/test_installation.py +++ b/tests/integration/install/test_installation.py @@ -1,6 +1,7 @@ import dataclasses import json import logging +import time from datetime import timedelta from typing import NoReturn @@ -19,6 +20,7 @@ AlreadyExists, InvalidParameterValue, NotFound, + ResourceDoesNotExist, ) from databricks.sdk.retries import retried @@ -393,8 +395,18 @@ def test_check_inventory_database_exists(ws, installation_ctx): def test_compare_remote_local_install_versions(ws, installation_ctx): - installation_finished = installation_ctx.workspace_installation.run() - assert installation_finished + installation_ctx.workspace_installation.run() + + # max time to wait for the installation to finish + timeout_duration = 5 + + 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) with pytest.raises( RuntimeWarning,