diff --git a/tests/integration/helpers/cluster.py b/tests/integration/helpers/cluster.py index 548b58a17e8..3914192f9af 100644 --- a/tests/integration/helpers/cluster.py +++ b/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: logginnfo( + 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() diff --git a/tests/integration/helpers/retry_decorator.py b/tests/integration/helpers/retry_decorator.py new file mode 100644 index 00000000000..aaa040464c2 --- /dev/null +++ b/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