MT-59 Refactor the CLI to use plain argparse (#52)

* Add regression tests for command --help

Signed-off-by: Fabrice Normandin <normandf@mila.quebec>

* Convert the CLI to plain argparse

- Adds unit tests to establish the current CLI behaviour in regression files:
  - "--help" output
  - ouput when running invalid commands
  - output when running commands that don't involve ssh'ing to the Mila cluster
- Refactors the CLI to be in plain argparse
  - Uglier for sure, but more transparent, easier to customize and extend to
    support more arguments (e.g. --cluster=narval option for mila code, etc)
  - Remove coleo as a dependency
  -  Add typing_extensions as a dependency
- Updates the regression files to reflect changes:
  - Before running mila or mila serve without choosing a subcommand would
    print nothing. Subcommands are now required.

Signed-off-by: Fabrice Normandin <normandf@mila.quebec>

* Apply suggestions from code review

Co-authored-by: Olivier Breuleux <breuleux@gmail.com>

* Remove commented lines

Signed-off-by: Fabrice Normandin <normandf@mila.quebec>

* Pass missing "port" argument down to _forward

Signed-off-by: Fabrice Normandin <normandf@mila.quebec>

* Fix bug in `mila serve tensorboard`

Signed-off-by: Fabrice Normandin <normandf@mila.quebec>

* Fix the docstring of test

Signed-off-by: Fabrice Normandin <normandf@mila.quebec>

---------

Signed-off-by: Fabrice Normandin <normandf@mila.quebec>
Co-authored-by: Olivier Breuleux <breuleux@gmail.com>

milatools/cli/remote.py

@@ -187,7 +187,8 @@ def home(self):
 
     def persist(self):
         qn.print(
-            "Warning: --persist does not work with --node or --job", style="orange"
+            "Warning: --persist does not work with --node or --job",
+            style="orange",
         )
         return self
 

pyproject.toml

@@ -11,10 +11,10 @@ license = "MIT"
 [tool.poetry.dependencies]
 python = "^3.7"
 blessed = "^1.18.1"
-coleo = "^0.3.0"
 sshconf = "^0.2.2"
 questionary = "^1.10.0"
 Fabric = "^2.7.0"
+typing-extensions = "^4.7.1"
 
 [tool.poetry.dev-dependencies]
 black = ">= 21.8b0"

tests/cli/test_commands.py

@@ -0,0 +1,109 @@
+import contextlib
+import io
+import os
+import shlex
+
+import pytest
+from pytest_regressions.file_regression import FileRegressionFixture
+
+from milatools.cli.commands import main
+
+
+def _convert_argparse_output_to_pre_py311_format(output: str) -> str:
+    """Revert the slight change in the output format of argparse in python 3.11."""
+    return output.replace("options:\n", "optional arguments:\n")
+
+
+@pytest.mark.parametrize(
+    "command",
+    ["mila"]
+    + [
+        f"mila {command}"
+        for command in ["docs", "intranet", "init", "forward", "code", "serve"]
+    ]
+    + [
+        f"mila serve {serve_subcommand}"
+        for serve_subcommand in (
+            "connect",
+            "kill",
+            "list",
+            "lab",
+            "notebook",
+            "tensorboard",
+            "mlflow",
+            "aim",
+        )
+    ],
+)
+def test_help(
+    command: str,
+    file_regression: FileRegressionFixture,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """Test that the --help text matches what's expected (and is stable over time)."""
+    monkeypatch.setattr("sys.argv", shlex.split(command + " --help"))
+    buf = io.StringIO()
+    with contextlib.suppress(SystemExit), contextlib.redirect_stdout(
+        buf
+    ), contextlib.redirect_stderr(buf):
+        with pytest.raises(SystemExit):
+            main()
+
+    output: str = buf.getvalue()
+    file_regression.check(_convert_argparse_output_to_pre_py311_format(output))
+
+
+@pytest.mark.parametrize(
+    "command",
+    [
+        "mila",  # Error: Missing a subcommand.
+        "mila search conda",
+        "mila code",  # Error: Missing the required PATH argument.
+        "mila serve",  # Error: Missing the subcommand.
+        "mila forward",  # Error: Missing the REMOTE argument.
+    ],
+)
+def test_invalid_command_output(
+    command: str,
+    file_regression: FileRegressionFixture,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """Test that we get a proper output when we use an invalid command (that exits immediately)."""
+    monkeypatch.setattr("sys.argv", shlex.split(command))
+    buf = io.StringIO()
+    with contextlib.suppress(SystemExit), pytest.raises(
+        SystemExit
+    ), contextlib.redirect_stdout(buf), contextlib.redirect_stderr(buf):
+        main()
+    file_regression.check(_convert_argparse_output_to_pre_py311_format(buf.getvalue()))
+
+
+# TODO: Perhaps we could use something like this so we can run all tests locally, but skip the ones
+# that need to actually connect to the cluster when running on GitHub Actions.
+# def dont_run_on_github(*args):
+#     return pytest.param(
+#         *args,
+#         marks=pytest.mark.skipif(
+#             "GITHUB_ACTIONS" in os.environ,
+#             reason="We don't run this test on GitHub Actions for security reasons.",
+#         ),
+#     )
+
+
+@pytest.mark.parametrize(
+    "command", ["mila docs conda", "mila intranet", "mila intranet idt"]
+)
+def test_check_command_output(
+    command: str,
+    file_regression: FileRegressionFixture,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """Run simple commands and check that their output matches what's expected."""
+
+    monkeypatch.setattr("webbrowser.open", lambda url: None)
+    monkeypatch.setattr("sys.argv", shlex.split(command))
+    buf = io.StringIO()
+    with contextlib.redirect_stdout(buf), contextlib.redirect_stderr(buf):
+        main()
+    output: str = buf.getvalue()
+    file_regression.check(_convert_argparse_output_to_pre_py311_format(output))

tests/cli/test_commands/test_check_command_output_mila_docs_conda_.txt

@@ -0,0 +1 @@
+Opening the docs: https://docs.mila.quebec/search.html?q=conda

tests/cli/test_commands/test_check_command_output_mila_intranet_.txt

@@ -0,0 +1 @@
+Opening the intranet: https://intranet.mila.quebec

tests/cli/test_commands/test_check_command_output_mila_intranet_idt_.txt

@@ -0,0 +1 @@
+Opening the intranet: https://sites.google.com/search/mila.quebec/mila-intranet?query=idt&scope=site&showTabs=false

tests/cli/test_commands/test_help_mila_.txt

@@ -0,0 +1,19 @@
+usage: mila [-h] [--version] {docs,intranet,init,forward,code,serve} ...
+
+Tools to connect to and interact with the Mila cluster. Cluster documentation:
+https://docs.mila.quebec/
+
+positional arguments:
+  {docs,intranet,init,forward,code,serve}
+    docs                Open the Mila cluster documentation.
+    intranet            Open the Mila intranet in a browser.
+    init                Set up your configuration and credentials.
+    forward             Forward a port on a compute node to your local
+                        machine.
+    code                Open a remote VSCode session on a compute node.
+    serve               Start services on compute nodes and forward them to
+                        your local machine.
+
+optional arguments:
+  -h, --help            show this help message and exit
+  --version, -v         Milatools version

tests/cli/test_commands/test_help_mila_code_.txt

@@ -0,0 +1,15 @@
+usage: mila code [-h] [--alloc ...] [--command VALUE] [--job VALUE]
+                 [--node VALUE] [--persist]
+                 PATH
+
+positional arguments:
+  PATH             Path to open on the remote machine
+
+optional arguments:
+  -h, --help       show this help message and exit
+  --alloc ...      Extra options to pass to slurm
+  --command VALUE  Command to use to start vscode (defaults to "code" or the
+                   value of $MILATOOLS_CODE_COMMAND)
+  --job VALUE      Job ID to connect to
+  --node VALUE     Node to connect to
+  --persist        Whether the server should persist or not

tests/cli/test_commands/test_help_mila_docs_.txt

@@ -0,0 +1,7 @@
+usage: mila docs [-h] ...
+
+positional arguments:
+  SEARCH      Search terms
+
+optional arguments:
+  -h, --help  show this help message and exit

tests/cli/test_commands/test_help_mila_forward_.txt

@@ -0,0 +1,9 @@
+usage: mila forward [-h] [--page VALUE] [--port VALUE] REMOTE
+
+positional arguments:
+  REMOTE        node:port to forward
+
+optional arguments:
+  -h, --help    show this help message and exit
+  --page VALUE  String to append after the URL
+  --port VALUE  Port to open on the local machine

tests/cli/test_commands/test_help_mila_init_.txt

@@ -0,0 +1,4 @@
+usage: mila init [-h]
+
+optional arguments:
+  -h, --help  show this help message and exit