Fix test for error output, make commands required

Signed-off-by: Fabrice Normandin <[email protected]>
mila-iqia · Aug 28, 2023 · ba8ed19 · ba8ed19
1 parent fe0cca2
commit ba8ed19
Show file tree

Hide file tree

Showing 7 changed files with 20 additions and 6 deletions.
diff --git a/milatools/cli/commands.py b/milatools/cli/commands.py
@@ -108,7 +108,7 @@ def mila():
         help="Milatools version",
     )
 
-    subparsers = parser.add_subparsers()
+    subparsers = parser.add_subparsers(required=True, dest="<command>")
 
     # ----- mila docs ------
 
@@ -149,7 +149,6 @@ def mila():
     forward_parser.add_argument("REMOTE", help="node:port to forward")
     forward_parser.add_argument(
         "--page",
-        # nargs="?",
         help="String to append after the URL",
         default=None,
         metavar="VALUE",
@@ -216,7 +215,9 @@ def mila():
         help="Start services on compute nodes and forward them to your local machine.",
         formatter_class=SortingHelpFormatter,
     )
-    serve_subparsers = serve_parser.add_subparsers(dest="serve_command", required=True)
+    serve_subparsers = serve_parser.add_subparsers(
+        dest="<serve_subcommand>", required=True
+    )
 
     # ----- mila serve connect ------
 
@@ -228,7 +229,6 @@ def mila():
     serve_connect_parser.add_argument(
         "IDENTIFIER",
         type=str,
-        # todo:remove spaces
         help="Server identifier output by the original mila serve command",
     )
     serve_connect_parser.add_argument(
@@ -345,7 +345,8 @@ def mila():
     args = parser.parse_args()
     args_dict = vars(args)
     function = args_dict.pop("function")
-    _ = args_dict.pop("serve_command", None)
+    _ = args_dict.pop("<command>")
+    _ = args_dict.pop("<serve_subcommand>", None)
     # replace SEARCH -> "search", REMOTE -> "remote", etc.
     args_dict = _convert_uppercase_keys_to_lowercase(args_dict)
     assert callable(function)

diff --git a/tests/cli/test_commands.py b/tests/cli/test_commands.py
@@ -71,7 +71,7 @@ def test_invalid_command_output(
     buf = io.StringIO()
     with contextlib.suppress(SystemExit), pytest.raises(
         SystemExit
-    ), contextlib.redirect_stdout(buf):
+    ), contextlib.redirect_stdout(buf), contextlib.redirect_stderr(buf):
         main()
     file_regression.check(_convert_argparse_output_to_pre_py311_format(buf.getvalue()))
 

diff --git a/tests/cli/test_commands/test_invalid_command_output_mila_.txt b/tests/cli/test_commands/test_invalid_command_output_mila_.txt
@@ -0,0 +1,2 @@
+usage: mila [-h] [--version] {docs,intranet,init,forward,code,serve} ...
+mila: error: the following arguments are required: <command>
diff --git a/tests/cli/test_commands/test_invalid_command_output_mila_code_.txt b/tests/cli/test_commands/test_invalid_command_output_mila_code_.txt
@@ -0,0 +1,4 @@
+usage: mila code [-h] [--alloc ...] [--command VALUE] [--job VALUE]
+                 [--node VALUE] [--persist]
+                 PATH
+mila code: error: the following arguments are required: PATH
diff --git a/tests/cli/test_commands/test_invalid_command_output_mila_forward_.txt b/tests/cli/test_commands/test_invalid_command_output_mila_forward_.txt
@@ -0,0 +1,2 @@
+usage: mila forward [-h] [--page VALUE] [--port VALUE] REMOTE
+mila forward: error: the following arguments are required: REMOTE
diff --git a/tests/cli/test_commands/test_invalid_command_output_mila_search_conda_.txt b/tests/cli/test_commands/test_invalid_command_output_mila_search_conda_.txt
@@ -0,0 +1,2 @@
+usage: mila [-h] [--version] {docs,intranet,init,forward,code,serve} ...
+mila: error: argument <command>: invalid choice: 'search' (choose from 'docs', 'intranet', 'init', 'forward', 'code', 'serve')
diff --git a/tests/cli/test_commands/test_invalid_command_output_mila_serve_.txt b/tests/cli/test_commands/test_invalid_command_output_mila_serve_.txt
@@ -0,0 +1,3 @@
+usage: mila serve [-h]
+                  {connect,kill,list,lab,notebook,tensorboard,mlflow,aim} ...
+mila serve: error: the following arguments are required: <serve_subcommand>