From f26664a569c026f2b4821e06cf7f747bd26df42b Mon Sep 17 00:00:00 2001
From: Peter Rowlands <peter@pmrowla.com>
Date: Tue, 24 May 2022 15:32:12 +0900
Subject: [PATCH] exp: re-enable run --run-all with celery queue

---
 dvc/repo/experiments/__init__.py           | 1 +
 dvc/repo/experiments/run.py                | 3 ++-
 tests/func/experiments/test_experiments.py | 8 ++++----
 3 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py
index 5d507fe71..69c88d4ec 100644
--- a/dvc/repo/experiments/__init__.py
+++ b/dvc/repo/experiments/__init__.py
@@ -168,6 +168,7 @@ class Experiments:
         if not entries:
             return results
 
+        # TODO: re-enable --jobs concurrency
         self.celery_queue.spawn_worker()
         failed = []
         for entry in entries:
diff --git a/dvc/repo/experiments/run.py b/dvc/repo/experiments/run.py
index ea042a216..8b63adc6b 100644
--- a/dvc/repo/experiments/run.py
+++ b/dvc/repo/experiments/run.py
@@ -25,7 +25,8 @@ def run(
     of `repro` for that experiment.
     """
     if run_all:
-        return repo.experiments.reproduce_queued(jobs=jobs)
+        entries = list(repo.experiments.celery_queue.iter_queued())
+        return repo.experiments.reproduce_celery(entries, jobs=jobs)
 
     if params:
         params = loads_param_overrides(params)
diff --git a/tests/func/experiments/test_experiments.py b/tests/func/experiments/test_experiments.py
index 9d218818a..756d504e4 100644
--- a/tests/func/experiments/test_experiments.py
+++ b/tests/func/experiments/test_experiments.py
@@ -577,16 +577,16 @@ def test_subrepo(tmp_dir, scm, workspace):
     assert resolve_rev(scm, ref_info.name) == exp
 
 
-@pytest.mark.xfail(reason="celery tests disabled")
-def test_queue(tmp_dir, scm, dvc, exp_stage, mocker):
+def test_run_celery(tmp_dir, scm, dvc, exp_stage, mocker):
+    """Test running with full (non-pytest-celery) dvc-task queue."""
     dvc.experiments.run(exp_stage.addressing, params=["foo=2"], queue=True)
     dvc.experiments.run(exp_stage.addressing, params=["foo=3"], queue=True)
     assert len(dvc.experiments.stash_revs) == 2
 
-    repro_mock = mocker.spy(dvc.experiments, "_reproduce_revs")
+    repro_spy = mocker.spy(dvc.experiments, "reproduce_celery")
     results = dvc.experiments.run(run_all=True)
     assert len(results) == 2
-    repro_mock.assert_called_with(jobs=1)
+    repro_spy.assert_called_once_with(entries=mocker.ANY, jobs=1)
 
     expected = {"foo: 2", "foo: 3"}
     metrics = set()