diff --git a/conftest.py b/conftest.py index 212b6a7c1..f2003dc0c 100644 --- a/conftest.py +++ b/conftest.py @@ -12,6 +12,21 @@ from daras_ai_v2.send_email import pytest_outbox +def flaky(fn): + max_tries = 5 + + @wraps(fn) + def wrapper(*args, **kwargs): + for i in range(max_tries): + try: + return fn(*args, **kwargs) + except Exception: + if i == max_tries - 1: + raise + + return wrapper + + @pytest.fixture(scope="session") def django_db_setup(django_db_setup, django_db_blocker): with django_db_blocker.unblock(): diff --git a/tests/test_translation.py b/tests/test_translation.py index d06b56d63..9ab1c23c0 100644 --- a/tests/test_translation.py +++ b/tests/test_translation.py @@ -1,3 +1,4 @@ +from conftest import flaky from daras_ai_v2.asr import run_google_translate @@ -48,6 +49,7 @@ def test_run_google_translate(threadpool_subtest): threadpool_subtest(_test_run_google_translate_one, text, expected) +@flaky def _test_run_google_translate_one( text: str, expected: str, glossary_url=None, target_lang="en" ):