From 8cbc4c43f154fb6505a4cc852d74d42e1ea0a4e4 Mon Sep 17 00:00:00 2001
From: Peter Rowlands <peter@pmrowla.com>
Date: Thu, 9 Mar 2023 14:36:16 +0900
Subject: [PATCH] tests: use session-scoped exp celery fixtures when possible

---
 tests/func/experiments/conftest.py               | 5 ++++-
 tests/func/experiments/test_apply.py             | 2 +-
 tests/func/experiments/test_experiments.py       | 2 +-
 tests/func/experiments/test_queue.py             | 2 ++
 tests/unit/repo/experiments/queue/test_celery.py | 6 +++---
 5 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/tests/func/experiments/conftest.py b/tests/func/experiments/conftest.py
index 9bc19b89c..7358dc098 100644
--- a/tests/func/experiments/conftest.py
+++ b/tests/func/experiments/conftest.py
@@ -5,6 +5,9 @@ from tests.unit.repo.experiments.conftest import (  # noqa, pylint disable=unuse
     exp_stage,
     failed_checkpoint_stage,
     failed_exp_stage,
+    session_app,
+    session_queue,
+    session_worker,
     test_queue,
 )
 
@@ -24,7 +27,7 @@ def http_auth_patch(mocker):
 
 
 @pytest.fixture(params=[True, False])
-def workspace(request, test_queue) -> bool:  # noqa: F811
+def workspace(request, session_queue) -> bool:  # noqa: F811
     return request.param
 
 
diff --git a/tests/func/experiments/test_apply.py b/tests/func/experiments/test_apply.py
index db3dcec4c..fe4e24663 100644
--- a/tests/func/experiments/test_apply.py
+++ b/tests/func/experiments/test_apply.py
@@ -26,7 +26,7 @@ def test_apply(tmp_dir, scm, dvc, exp_stage):
     assert (tmp_dir / "metrics.yaml").read_text().strip() == "foo: 3"
 
 
-def test_apply_failed(tmp_dir, scm, dvc, failed_exp_stage, test_queue):
+def test_apply_failed(tmp_dir, scm, dvc, failed_exp_stage, session_queue):
     from dvc.exceptions import InvalidArgumentError
 
     dvc.experiments.run(failed_exp_stage.addressing, params=["foo=2"], queue=True)
diff --git a/tests/func/experiments/test_experiments.py b/tests/func/experiments/test_experiments.py
index 1a5b32543..4a1bbe7fa 100644
--- a/tests/func/experiments/test_experiments.py
+++ b/tests/func/experiments/test_experiments.py
@@ -131,7 +131,7 @@ def test_get_baseline(tmp_dir, scm, dvc, exp_stage):
     assert dvc.experiments.get_baseline(f"{CELERY_STASH}@{{1}}") == init_rev
 
 
-def test_update_py_params(tmp_dir, scm, dvc, test_queue, copy_script):
+def test_update_py_params(tmp_dir, scm, dvc, session_queue, copy_script):
     tmp_dir.gen("params.py", "INT = 1\n")
     stage = dvc.run(
         cmd="python copy.py params.py metrics.py",
diff --git a/tests/func/experiments/test_queue.py b/tests/func/experiments/test_queue.py
index 0aac40706..51c685733 100644
--- a/tests/func/experiments/test_queue.py
+++ b/tests/func/experiments/test_queue.py
@@ -19,6 +19,7 @@ def test_celery_logs(
     failed_exp_stage,
     follow,
     capsys,
+    test_queue,
 ):
     celery_queue = dvc.experiments.celery_queue
     dvc.experiments.run(failed_exp_stage.addressing, queue=True)
@@ -37,6 +38,7 @@ def test_queue_remove_done(
     dvc,
     exp_stage,
     failed_exp_stage,
+    test_queue,
 ):
     queue_length = 3
     success_tasks = []
diff --git a/tests/unit/repo/experiments/queue/test_celery.py b/tests/unit/repo/experiments/queue/test_celery.py
index 42c0ad490..0490d5804 100644
--- a/tests/unit/repo/experiments/queue/test_celery.py
+++ b/tests/unit/repo/experiments/queue/test_celery.py
@@ -49,10 +49,10 @@ def test_shutdown_with_kill(test_queue, mocker):
     )
 
 
-def test_post_run_after_kill(test_queue):
+def test_post_run_after_kill(session_queue):
     from celery import chain
 
-    sig_bar = test_queue.proc.run_signature(
+    sig_bar = session_queue.proc.run_signature(
         ["python3", "-c", "import time; time.sleep(10)"], name="bar"
     )
     sig_bar.freeze()
@@ -65,7 +65,7 @@ def test_post_run_after_kill(test_queue):
 
     while True:
         try:
-            test_queue.proc.kill("bar")
+            session_queue.proc.kill("bar")
             assert result_foo.status == "PENDING"
             break
         except ProcessLookupError: