Migrate service-status from argparse to click (#48)

* Migrate service-status from argparse to click * Update obs_common/service_status.py Co-authored-by: Will Kahn-Greene <[email protected]> --------- Co-authored-by: Will Kahn-Greene <[email protected]>
mozilla-services · Nov 7, 2024 · 3796458 · 3796458
1 parent 213fedd
commit 3796458
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 82 deletions.
diff --git a/obs_common/service_status.py b/obs_common/service_status.py
@@ -9,63 +9,22 @@
 how far behind different server environments are from main tip.
 """
 
-import argparse
 import json
-import os
 import sys
 from urllib.parse import urlparse
 from urllib.request import urlopen
+from pathlib import Path
 
+import click
+import tomllib
 
-DESCRIPTION = """
-service-status.py tells you how far behind different server environments
-are from main tip.
-"""
-
-DEFAULT_CONFIG = {
-    # The name of the main branch in the repository
-    "main_branch": "main",
-    # List of "label=host" for hosts that have a /__version__ to check
-    "hosts": [],
-}
-
-
-def get_config():
-    """Generates configuration.
-
-    This tries to pull configuration from the ``[tool.service-status]`` table
-    from a ``pyproject.toml`` file.
-
-    If neither exist, then it uses defaults.
-
-    :returns: configuration dict
-
-    """
-    my_config = dict(DEFAULT_CONFIG)
-
-    if os.path.exists("pyproject.toml"):
-        if sys.version_info >= (3, 11):
-            import tomllib
-        else:
-            try:
-                import tomli as tomllib
-            except ImportError:
-                print(
-                    "For Python <3.11, you need to install tomli to work with pyproject.toml "
-                    + "files."
-                )
-                tomllib = None
-
-        if tomllib is not None:
-            with open("pyproject.toml", "rb") as fp:
-                data = tomllib.load(fp)
 
-            config_data = data.get("tool", {}).get("service-status", {})
-            if config_data:
-                for key, default_val in my_config.items():
-                    my_config[key] = config_data.get(key, default_val)
+DESCRIPTION = """
+Report how far behind different server environments are from main tip.
 
-    return my_config
+For options that are not specified, values are pulled from the [tool.service-status]
+section in the pyproject.toml file in the current working directory, if it exists.
+"""
 
 
 def fetch(url, is_json=True):
@@ -141,44 +100,39 @@ def print_delta(self, main_branch, user, repo, sha):
         self.row()
 
 
-def main():
-    config = get_config()
-
-    parser = argparse.ArgumentParser(description=DESCRIPTION)
+@click.command(help=DESCRIPTION)
+@click.option(
+    "--main-branch",
+    help=(
+        "The name of the main branch in the git repository. Defaults "
+        'to "main" if not configured in pyproject.toml.'
+    ),
+)
+@click.option(
+    "--host",
+    "hosts",
+    multiple=True,
+    help=(
+        "A list of hosts which have a ``/__version__`` Dockerflow endpoint in the "
+        "form of ENVIRONMENTNAME=HOST."
+    ),
+)
+def main(main_branch, hosts):
+    config_data = {}
+    if (pyproject_toml := Path("pyproject.toml")).exists():
+        data = tomllib.loads(pyproject_toml.read_text())
+        config_data = data.get("tool", {}).get("service-status", {})
+
+    main_branch = main_branch or config_data.get("main_branch", "main")
 
-    # Add items that can be configured to argparse as configuration options.
-    # This makes it possible to specify or override configuration with command
-    # line arguments.
-    for key, val in config.items():
-        key_arg = key.replace("_", "-")
-        if isinstance(val, list):
-            parser.add_argument(
-                f"--{key_arg}",
-                default=val,
-                nargs="+",
-                metavar="VALUE",
-                help=f"override configuration {key}; defaults to {val!r}",
-            )
+    if not hosts:
+        if "hosts" in config_data:
+            hosts = config_data["hosts"]
         else:
-            default_val = val.replace("%", "%%")
-            parser.add_argument(
-                f"--{key_arg}",
-                default=val,
-                metavar="VALUE",
-                help=f"override configuration {key}; defaults to {default_val!r}",
-            )
-
-    args = parser.parse_args()
-
-    main_branch = args.main_branch
-    hosts = args.hosts
+            raise click.ClickException("no hosts configured")
 
     out = StdoutOutput()
 
-    if not hosts:
-        print("no hosts specified.")
-        return 1
-
     current_section = ""
 
     for line in hosts:

diff --git a/tests/test_service_status.py b/tests/test_service_status.py
@@ -0,0 +1,14 @@
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+from click.testing import CliRunner
+
+from obs_common import service_status
+
+
+def test_it_runs():
+    """Test whether the module loads and spits out help."""
+    runner = CliRunner()
+    result = runner.invoke(service_status.main, ["--help"])
+    assert result.exit_code == 0