From f43c04f93537f51b96edb205e7e31c1b7b69086c Mon Sep 17 00:00:00 2001 From: karajan1001 Date: Sun, 24 Oct 2021 17:56:24 +0800 Subject: [PATCH] Make machine status name optional fix: #6838 1. make machine status name optional. 2. show all machine status if no name provide. 3. UI changes for this cmd. 4. add a new test for it. Co-authored-by: Dave Berenbaum --- dvc/command/machine.py | 34 +++++++++++++++++------ tests/func/machine/test_machine_status.py | 5 ++-- tests/unit/command/test_machine.py | 18 ++++++++++-- 3 files changed, 44 insertions(+), 13 deletions(-) diff --git a/dvc/command/machine.py b/dvc/command/machine.py index a6541ef40e..2de8d48f92 100644 --- a/dvc/command/machine.py +++ b/dvc/command/machine.py @@ -202,16 +202,32 @@ class CmdMachineStatus(CmdBase): "instance_gpu", ] + def _show_machine_status(self, name: str): + ui.write(f"machine '{name}':") + all_status = list(self.repo.machine.status(name)) + if not all_status: + ui.write("\toffline") + for i, status in enumerate(all_status, start=1): + ui.write(f"\tinstance_num_{i}:") + for field in self.SHOWN_FIELD: + value = status.get(field, None) + ui.write(f"\t\t{field:20}: {value}") + def run(self): if self.repo.machine is None: raise MachineDisabledError - all_status = self.repo.machine.status(self.args.name) - for i, status in enumerate(all_status, start=1): - ui.write(f"instance_num_{i}:") - for field in self.SHOWN_FIELD: - value = status.get(field, None) - ui.write(f"\t{field:20}: {value}") + if self.args.name: + self._show_machine_status(self.args.name) + else: + name_set = set() + for level in self.repo.config.LEVELS: + conf = self.repo.config.read(level)["machine"] + name_set.update(conf.keys()) + name_list = list(name_set) + for name in sorted(name_list): + self._show_machine_status(name) + return 0 @@ -390,7 +406,9 @@ def add_parser(subparsers, parent_parser): ) machine_create_parser.set_defaults(func=CmdMachineCreate) - machine_STATUS_HELP = "List the status of a running machine." + machine_STATUS_HELP = ( + "List the status of running instances for one/all machines." + ) machine_status_parser = machine_subparsers.add_parser( "status", parents=[parent_config_parser, parent_parser], @@ -399,7 +417,7 @@ def add_parser(subparsers, parent_parser): formatter_class=argparse.RawDescriptionHelpFormatter, ) machine_status_parser.add_argument( - "name", help="Name of the running machine." + "name", nargs="?", help="(optional) Name of the machine." ) machine_status_parser.set_defaults(func=CmdMachineStatus) diff --git a/tests/func/machine/test_machine_status.py b/tests/func/machine/test_machine_status.py index 9bea45e558..33b474be78 100644 --- a/tests/func/machine/test_machine_status.py +++ b/tests/func/machine/test_machine_status.py @@ -8,6 +8,7 @@ def test_status( status = machine_instance assert main(["machine", "status", "foo"]) == 0 cap = capsys.readouterr() - assert "instance_num_1:" in cap.out + assert "machine 'foo':\n" in cap.out + assert "\tinstance_num_1:\n" in cap.out for key in CmdMachineStatus.SHOWN_FIELD: - assert f"\t{key:20}: {status[key]}" in cap.out + assert f"\t\t{key:20}: {status[key]}\n" in cap.out diff --git a/tests/unit/command/test_machine.py b/tests/unit/command/test_machine.py index 4177a89aba..049db10222 100644 --- a/tests/unit/command/test_machine.py +++ b/tests/unit/command/test_machine.py @@ -1,4 +1,5 @@ import configobj +from mock import call from dvc.cli import parse_args from dvc.command.machine import ( @@ -19,6 +20,8 @@ "[feature]\n" " machine = true\n" "['machine \"foo\"']\n" + " cloud = aws\n" + "['machine \"myaws\"']\n" " cloud = aws" ) } @@ -42,7 +45,7 @@ def test_remove(tmp_dir): cmd = cli_args.func(cli_args) assert cmd.run() == 0 config = configobj.ConfigObj(str(tmp_dir / ".dvc" / "config")) - assert list(config.keys()) == ["feature"] + assert list(config.keys()) == ["feature", 'machine "myaws"'] def test_create(tmp_dir, dvc, mocker): @@ -58,7 +61,8 @@ def test_create(tmp_dir, dvc, mocker): m.assert_called_once_with("foo") -def test_status(tmp_dir, dvc, mocker): +def test_status(tmp_dir, scm, dvc, mocker): + tmp_dir.gen(DATA) cli_args = parse_args(["machine", "status", "foo"]) assert cli_args.func == CmdMachineStatus @@ -66,10 +70,18 @@ def test_status(tmp_dir, dvc, mocker): m = mocker.patch.object( cmd.repo.machine, "status", autospec=True, return_value=[] ) - assert cmd.run() == 0 m.assert_called_once_with("foo") + cli_args = parse_args(["machine", "status"]) + cmd = cli_args.func(cli_args) + m = mocker.patch.object( + cmd.repo.machine, "status", autospec=True, return_value=[] + ) + assert cmd.run() == 0 + assert m.call_count == 2 + m.assert_has_calls([call("foo"), call("myaws")]) + def test_destroy(tmp_dir, dvc, mocker): cli_args = parse_args(["machine", "destroy", "foo"])