exp queue: Add two command queue start and stop (#7589)

fix: #7589

1. Add two sub-command `dvc queue start` and `dvc queue stop`
2. Add a unit test to test them
This commit is contained in:
karajan1001
2022-07-12 21:39:49 +09:00
committed by Peter Rowlands (변기호)
parent a6cb4eb72e
commit 12fc8c4e7d
4 changed files with 137 additions and 1 deletions
+34
View File
@@ -1,6 +1,8 @@
from dvc.cli import parse_args
from dvc.commands.queue.kill import CmdQueueKill
from dvc.commands.queue.remove import CmdQueueRemove
from dvc.commands.queue.start import CmdQueueStart
from dvc.commands.queue.stop import CmdQueueStop
def test_experiments_remove(dvc, scm, mocker):
@@ -61,3 +63,35 @@ def test_experiments_kill(dvc, scm, mocker):
assert cmd.run() == 0
m.assert_called_once_with(revs=["exp1", "exp2"])
def test_experiments_start(dvc, scm, mocker):
cli_args = parse_args(["queue", "start", "-j", "3"])
assert cli_args.func == CmdQueueStart
cmd = cli_args.func(cli_args)
m = mocker.patch(
"dvc.repo.experiments.queue.local.LocalCeleryQueue.spawn_worker",
)
assert cmd.run() == 0
assert m.call_count == 3
def test_experiments_stop(dvc, scm, mocker):
cli_args = parse_args(
[
"queue",
"stop",
"--kill",
]
)
assert cli_args.func == CmdQueueStop
cmd = cli_args.func(cli_args)
m = mocker.patch(
"dvc.repo.experiments.queue.local.LocalCeleryQueue.shutdown",
)
assert cmd.run() == 0
m.assert_called_once_with(kill=True)