From 56f90eb3c95ca19511d34f420988a1bef5cea82c Mon Sep 17 00:00:00 2001 From: Peter Rowlands Date: Thu, 5 May 2022 16:25:41 +0900 Subject: [PATCH] exp: add test_queue unit test helper for celery queue --- tests/unit/repo/experiments/conftest.py | 29 +++++++++++++++++++++++++ 1 file changed, 29 insertions(+) create mode 100644 tests/unit/repo/experiments/conftest.py diff --git a/tests/unit/repo/experiments/conftest.py b/tests/unit/repo/experiments/conftest.py new file mode 100644 index 000000000..9837c5f1c --- /dev/null +++ b/tests/unit/repo/experiments/conftest.py @@ -0,0 +1,29 @@ +from contextlib import contextmanager + +import pytest + +from dvc.repo.experiments.queue.local import LocalCeleryQueue + + +@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 + + +@pytest.fixture +def test_queue(tmp_dir, dvc, scm, mocker) -> LocalCeleryQueue: + """Patches experiments celery queue for pytest testing. + + 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, concurrency=1) as worker: + mocker.patch.object(celery_queue, "worker", return_value=worker) + yield celery_queue