Merge pull request #67019 from ClickHouse/divanik/fix_test_azure_flaky

[Green CI] Fix test test_storage_azure_blob_storage

tests/integration/helpers/cluster.py

@@ -52,6 +52,7 @@ from helpers.client import QueryRuntimeException
 import docker
 
 from .client import Client
+from .retry_decorator import retry
 
 from .config_cluster import *
 
@@ -2690,15 +2691,12 @@ class ClickHouseCluster:
 
             images_pull_cmd = self.base_cmd + ["pull"]
             # sometimes dockerhub/proxy can be flaky
-            for i in range(5):
-                try:
-                    run_and_check(images_pull_cmd)
-                    break
-                except Exception as ex:
-                    if i == 4:
-                        raise ex
-                    logging.info("Got exception pulling images: %s", ex)
-                    time.sleep(i * 3)
+
+            retry(
+                log_function=lambda exception: logging.info(
+                    "Got exception pulling images: %s", exception
+                ),
+            )(run_and_check)(images_pull_cmd)
 
             if self.with_zookeeper_secure and self.base_zookeeper_cmd:
                 logging.debug("Setup ZooKeeper Secure")
@@ -2971,7 +2969,11 @@ class ClickHouseCluster:
                     "Trying to create Azurite instance by command %s",
                     " ".join(map(str, azurite_start_cmd)),
                 )
-                run_and_check(azurite_start_cmd)
+                retry(
+                    log_function=lambda exception: logging.info(
+                        f"Azurite initialization failed with error: {exception}"
+                    ),
+                )(run_and_check)(azurite_start_cmd)
                 self.up_called = True
                 logging.info("Trying to connect to Azurite")
                 self.wait_azurite_to_start()

tests/integration/helpers/retry_decorator.py

@@ -0,0 +1,36 @@
+import time
+import random
+from typing import Type, List
+
+
+def retry(
+    retries: int = 5,
+    delay: float = 1,
+    backoff: float = 1.5,
+    jitter: float = 2,
+    log_function=lambda *args, **kwargs: None,
+    retriable_expections_list: List[Type[BaseException]] = [Exception],
+):
+    def inner(func):
+        def wrapper(*args, **kwargs):
+            current_delay = delay
+            for retry in range(retries):
+                try:
+                    func(*args, **kwargs)
+                    break
+                except Exception as e:
+                    should_retry = False
+                    for retriable_exception in retriable_expections_list:
+                        if isinstance(e, retriable_exception):
+                            should_retry = True
+                            break
+                    if not should_retry or (retry == retries - 1):
+                        raise e
+                    log_function(retry=retry, exception=e)
+                    sleep_time = current_delay + random.uniform(0, jitter)
+                    time.sleep(sleep_time)
+                    current_delay *= backoff
+
+        return wrapper
+
+    return inner