tests: retry celery queue on timeout (#8488)

retry celery queue on timeout

tests/unit/repo/experiments/conftest.py

@@ -1,4 +1,4 @@
-from contextlib import contextmanager
+from functools import partial
 from textwrap import dedent
 
 import pytest
@@ -117,15 +117,13 @@ def failed_exp_stage(tmp_dir, scm, dvc):
     return stage
 
 
-@contextmanager
 def _thread_worker(app, **kwargs):
     # Based on pytest-celery's celery_worker fixture but using thread pool
     # instead of solo pool so that broadcast/control API is available
     from celery.contrib.testing import worker
 
     app.loader.import_task_module("celery.contrib.testing.tasks")
-    with worker.start_worker(app, pool="threads", **kwargs) as test_worker:
-        yield test_worker
+    return worker.start_worker(app, pool="threads", **kwargs)
 
 
 @pytest.fixture
@@ -134,12 +132,26 @@ def test_queue(tmp_dir, dvc, scm, mocker) -> LocalCeleryQueue:
 
     Test queue worker runs for the duration of the test in separate thread(s).
     """
-    celery_queue = dvc.experiments.celery_queue
-    mocker.patch.object(celery_queue, "spawn_worker")
-    with _thread_worker(
-        celery_queue.celery,
+    import celery
+
+    queue = dvc.experiments.celery_queue
+    mocker.patch.object(queue, "spawn_worker")
+
+    f = partial(
+        _thread_worker,
+        queue.celery,
         concurrency=1,
-        ping_task_timeout=30,
-    ) as worker:
-        mocker.patch.object(celery_queue, "worker", return_value=worker)
-        yield celery_queue
+        ping_task_timeout=20,
+    )
+    exc = None
+    for _ in range(3):
+        try:
+            with f() as worker:
+                mocker.patch.object(queue, "worker", return_value=worker)
+                yield queue
+                return
+        except celery.exceptions.TimeoutError as e:
+            exc = e
+            continue
+    assert exc
+    raise exc