def retry_flaky_test(tries=3):
    """Retries a flaky test a specified number of times."""

    def flaky_test_func(test_func):
        @wraps(test_func)
        def decorated_func(*args, **kwargs):
            for i in range(tries):
                try:
                    return test_func(*args, **kwargs)
                except Exception as e:
                    if i == tries - 1:
                        raise e
                    time.sleep(2)

        return decorated_func

    return flaky_test_func
Retries a flaky test a specified number of times.
ai/core/src/unitycatalog/ai/test_utils/client_utils.py