diff --git a/dvc/commands/experiments/remove.py b/dvc/commands/experiments/remove.py index 92ebf9bdf..bd7bf96e0 100644 --- a/dvc/commands/experiments/remove.py +++ b/dvc/commands/experiments/remove.py @@ -10,23 +10,22 @@ logger = logging.getLogger(__name__) class CmdExperimentsRemove(CmdBase): - def raise_error_if_all_disabled(self): + def argument_checking(self): if not any( [ - self.args.experiment, self.args.all_commits, self.args.rev, self.args.queue, ] - ): + ) ^ bool(self.args.experiment): raise InvalidArgumentError( "Either provide an `experiment` argument, or use the " - "`--rev` or `--all-commits` flag." + "`--rev` or `--all-commits` or `--queue` flag." ) def run(self): - self.raise_error_if_all_disabled() + self.argument_checking() removed_list = self.repo.experiments.remove( exp_names=self.args.experiment, diff --git a/tests/unit/command/test_experiments.py b/tests/unit/command/test_experiments.py index 506011803..49efe01fb 100644 --- a/tests/unit/command/test_experiments.py +++ b/tests/unit/command/test_experiments.py @@ -342,7 +342,7 @@ def test_experiments_pull(dvc, scm, mocker): ) -def test_experiments_remove(dvc, scm, mocker, capsys, caplog): +def test_experiments_remove_flag(dvc, scm, mocker, capsys, caplog): cli_args = parse_args( [ "experiments", @@ -354,19 +354,15 @@ def test_experiments_remove(dvc, scm, mocker, capsys, caplog): "2", "--git-remote", "myremote", - "exp-123", - "exp-234", ] ) assert cli_args.func == CmdExperimentsRemove - cmd = cli_args.func(cli_args) m = mocker.patch("dvc.repo.experiments.remove.remove", return_value={}) - assert cmd.run() == 0 m.assert_called_once_with( cmd.repo, - exp_names=["exp-123", "exp-234"], + exp_names=[], all_commits=True, rev="foo", num=2, @@ -374,12 +370,45 @@ def test_experiments_remove(dvc, scm, mocker, capsys, caplog): git_remote="myremote", ) - cmd = cli_args.func(parse_args(["exp", "remove"])) + +def test_experiments_remove_special(dvc, scm, mocker, capsys, caplog): + cli_args = parse_args( + [ + "experiments", + "remove", + "--git-remote", + "myremote", + "exp-123", + "exp-234", + ] + ) + assert cli_args.func == CmdExperimentsRemove + cmd = cli_args.func(cli_args) + m = mocker.patch("dvc.repo.experiments.remove.remove", return_value={}) + assert cmd.run() == 0 + m.assert_called_once_with( + cmd.repo, + exp_names=["exp-123", "exp-234"], + all_commits=False, + rev=None, + num=1, + queue=False, + git_remote="myremote", + ) + + +def test_experiments_remove_invalid(dvc, scm, mocker, capsys, caplog): + cmd = CmdExperimentsRemove( + parse_args(["exp", "remove", "--all-commits", "exp-1", "exp-2"]) + ) + with pytest.raises(InvalidArgumentError): + cmd.run() + cmd = CmdExperimentsRemove(parse_args(["exp", "remove"])) with pytest.raises(InvalidArgumentError) as excinfo: cmd.run() assert ( str(excinfo.value) == "Either provide an `experiment` argument" - ", or use the `--rev` or `--all-commits` flag." + ", or use the `--rev` or `--all-commits` or `--queue` flag." )