diff --git a/dvc/commands/queue/remove.py b/dvc/commands/queue/remove.py index c3b8bbe5b..760ab76ce 100644 --- a/dvc/commands/queue/remove.py +++ b/dvc/commands/queue/remove.py @@ -3,6 +3,7 @@ import logging from dvc.cli.command import CmdBase from dvc.cli.utils import append_doc_link +from dvc.exceptions import InvalidArgumentError from dvc.ui import ui logger = logging.getLogger(__name__) @@ -11,14 +12,40 @@ logger = logging.getLogger(__name__) class CmdQueueRemove(CmdBase): """Remove exp in queue.""" - def run(self): - removed_list = self.repo.experiments.celery_queue.remove( - revs=self.args.task, - all_=self.args.all, - success=self.args.success, - queued=self.args.queued, - failed=self.args.failed, + def argument_checking(self): + clear_flag = any( + [ + self.args.all, + self.args.queued, + self.args.failed, + self.args.success, + ] ) + if not (clear_flag ^ bool(self.args.task)): + raise InvalidArgumentError( + "Either provide an `tasks` argument, or use the " + "`--all`, `--queued`, `--failed`, `--success` flag." + ) + + def run(self): + + self.argument_checking() + + if self.args.all: + self.args.queued = True + self.args.failed = True + self.args.success = True + + if self.args.queued or self.args.failed or self.args.success: + removed_list = self.repo.experiments.celery_queue.clear( + success=self.args.success, + queued=self.args.queued, + failed=self.args.failed, + ) + else: + removed_list = self.repo.experiments.celery_queue.remove( + revs=self.args.task, + ) if removed_list: removed = ", ".join(removed_list) diff --git a/dvc/repo/experiments/queue/remove.py b/dvc/repo/experiments/queue/remove.py index 082666c31..830412863 100644 --- a/dvc/repo/experiments/queue/remove.py +++ b/dvc/repo/experiments/queue/remove.py @@ -14,7 +14,7 @@ from dvc.repo.experiments.queue.base import QueueDoneResult if TYPE_CHECKING: from dvc.repo.experiments.queue.base import QueueEntry - from dvc.repo.experiments.queue.local import LocalCeleryQueue + from dvc.repo.experiments.queue.celery import LocalCeleryQueue from dvc.repo.experiments.stash import ExpStashEntry @@ -134,27 +134,15 @@ def celery_clear(self: "LocalCeleryQueue", **kwargs) -> List[str]: def celery_remove( self: "LocalCeleryQueue", revs: Collection[str], - queued: bool = False, - failed: bool = False, - success: bool = False, - all_: bool = False, ) -> List[str]: """Remove the specified entries from the queue. Arguments: revs: Stash revisions or queued exp names to be removed. - queued: Remove all queued tasks. - failed: Remove all failed tasks. - success: Remove all success tasks. - all_: Remove all tasks. Returns: Revisions (or names) which were removed. """ - if all_: - queued = failed = success = True - if queued or failed or success: - return self.clear(failed=failed, success=success, queued=queued) # match_queued queue_match_results = self.match_queue_entry_by_name( diff --git a/tests/func/experiments/test_queue.py b/tests/func/experiments/test_queue.py index f610e0300..92259c31b 100644 --- a/tests/func/experiments/test_queue.py +++ b/tests/func/experiments/test_queue.py @@ -86,10 +86,7 @@ def test_queue_remove_done(dvc, failed_tasks, success_tasks): status = to_dict(dvc.experiments.celery_queue.status()) assert set(status) == set(failed_tasks[1:] + success_tasks[:2]) - assert ( - dvc.experiments.celery_queue.remove([], failed=True) - == failed_tasks[1:] - ) + assert dvc.experiments.celery_queue.clear(failed=True) == failed_tasks[1:] assert len(dvc.experiments.celery_queue.failed_stash) == 0 assert set(to_dict(dvc.experiments.celery_queue.status())) == set( @@ -97,8 +94,7 @@ def test_queue_remove_done(dvc, failed_tasks, success_tasks): ) assert ( - dvc.experiments.celery_queue.remove([], success=True) - == success_tasks[:2] + dvc.experiments.celery_queue.clear(success=True) == success_tasks[:2] ) assert dvc.experiments.celery_queue.status() == [] diff --git a/tests/unit/command/test_queue.py b/tests/unit/command/test_queue.py index 825200a2c..7f54e7952 100644 --- a/tests/unit/command/test_queue.py +++ b/tests/unit/command/test_queue.py @@ -7,36 +7,83 @@ from dvc.commands.queue.remove import CmdQueueRemove from dvc.commands.queue.start import CmdQueueStart from dvc.commands.queue.status import CmdQueueStatus from dvc.commands.queue.stop import CmdQueueStop +from dvc.exceptions import InvalidArgumentError -def test_experiments_remove(dvc, scm, mocker): +def test_experiments_remove_flags(dvc, scm, mocker): + cli_args = parse_args( + [ + "queue", + "remove", + "--queued", + "--success", + "--failed", + ] + ) + assert cli_args.func == CmdQueueRemove + cmd = cli_args.func(cli_args) + remove_mocker = mocker.patch( + "dvc.repo.experiments.queue.celery.LocalCeleryQueue.clear", + return_value={}, + ) + assert cmd.run() == 0 + remove_mocker.assert_called_once_with( + success=True, + failed=True, + queued=True, + ) cli_args = parse_args( [ "queue", "remove", "--all", - "--queued", - "--success", - "--failed", + ] + ) + cmd = cli_args.func(cli_args) + remove_mocker.reset_mock() + assert cmd.run() == 0 + remove_mocker.assert_called_once_with( + success=True, + failed=True, + queued=True, + ) + + +def test_experiments_remove_invalid(dvc, scm, mocker): + cli_args = parse_args(["queue", "remove", "--queued", ["exp1", "exp2"]]) + cmd = cli_args.func(cli_args) + with pytest.raises(InvalidArgumentError): + assert cmd.run() == 0 + + cli_args = parse_args( + [ + "queue", + "remove", + ] + ) + cmd = cli_args.func(cli_args) + with pytest.raises(InvalidArgumentError): + assert cmd.run() == 0 + + +def test_experiments_remove_name(dvc, scm, mocker): + cli_args = parse_args( + [ + "queue", + "remove", "exp1", "exp2", ] ) assert cli_args.func == CmdQueueRemove - cmd = cli_args.func(cli_args) remove_mocker = mocker.patch( "dvc.repo.experiments.queue.celery.LocalCeleryQueue.remove", return_value={}, ) - assert cmd.run() == 0 remove_mocker.assert_called_once_with( revs=["exp1", "exp2"], - all_=True, - success=True, - failed=True, - queued=True, ) diff --git a/tests/unit/repo/experiments/queue/test_remove.py b/tests/unit/repo/experiments/queue/test_remove.py index 2a0f1c5e6..8bb4165b2 100644 --- a/tests/unit/repo/experiments/queue/test_remove.py +++ b/tests/unit/repo/experiments/queue/test_remove.py @@ -37,7 +37,7 @@ def test_remove_queued(test_queue, mocker): remove_revs_mocker.reset_mock() reject_mocker.reset_mock() - assert test_queue.remove([], queued=True) == queued_test + assert test_queue.clear(queued=True) == queued_test remove_revs_mocker.assert_called_once_with(list(stash_dict.values())) reject_mocker.assert_has_calls( [call("msg_queue1"), call("msg_queue2"), call("msg_queue3")] @@ -103,7 +103,7 @@ def test_remove_done(test_queue, mocker): remove_revs_mocker.reset_mock() purge_mocker.reset_mock() - assert set(test_queue.remove([], success=True, failed=True)) == set( + assert set(test_queue.clear(success=True, failed=True)) == set( failed_test ) | set(success_test) purge_mocker.assert_has_calls( @@ -118,11 +118,3 @@ def test_remove_done(test_queue, mocker): any_order=True, ) remove_revs_mocker.assert_called_once_with(list(stash_dict.values())) - - -def test_remove_all(test_queue, mocker): - clear_mocker = mocker.patch.object( - test_queue, "clear", return_value=mocker.Mock() - ) - test_queue.remove([], all_=True) - assert clear_mocker.called_once_with(queud=True, failed=True, success=True)