diff --git a/tests/unit/repo/experiments/conftest.py b/tests/unit/repo/experiments/conftest.py index 3aba50586..a4968ca7d 100644 --- a/tests/unit/repo/experiments/conftest.py +++ b/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