From 4502f98e52d3c2c5376ed407fc06194f3659f7b3 Mon Sep 17 00:00:00 2001
From: Andreas Motl <andreas.motl@panodata.org>
Date: Tue, 12 Oct 2021 22:49:57 +0200
Subject: [PATCH] Add "gribmagic unity list" subcommand

---
 README.md                                     |  8 +++--
 gribmagic/commands.py                         | 21 ++++++++++--
 setup.py                                      |  2 +-
 .../{test_command.py => test_commands.py}     | 34 ++++++++++++++++---
 4 files changed, 55 insertions(+), 10 deletions(-)
 rename tests/unity/{test_command.py => test_commands.py} (78%)

diff --git a/README.md b/README.md
index 6cffcd2..4a60014 100644
--- a/README.md
+++ b/README.md
@@ -69,8 +69,12 @@ make test
 
 ### Ad hoc usage
 ```
+# List available labels.
+gribmagic unity list
+
+# Acquire data.
 export GM_DATA_PATH=.gribmagic-data
-gribmagic unity --model=dwd-icon-eu --timestamp=2021-10-03T00:00:00Z
+gribmagic unity acquire --model=dwd-icon-eu --timestamp=2021-10-03T00:00:00Z
 ```
 
 ### Configuration
@@ -96,7 +100,7 @@ docker build --tag gribmagic .
 
 and then invoke it like
 ```
-docker run -it --volume=$PWD/.gribmagic-data:/var/spool/gribmagic gribmagic:latest gribmagic unity --model=dwd-icon-eu --timestamp=2021-10-03T00:00:00Z
+docker run -it --volume=$PWD/.gribmagic-data:/var/spool/gribmagic gribmagic:latest gribmagic unity acquire --model=dwd-icon-eu --timestamp=2021-10-03T00:00:00Z
 ```
 
 ---
diff --git a/gribmagic/commands.py b/gribmagic/commands.py
index fd23d8d..2b25c7d 100644
--- a/gribmagic/commands.py
+++ b/gribmagic/commands.py
@@ -1,4 +1,5 @@
 """Command line entry points for NWP data download"""
+import json
 import logging
 from datetime import datetime
 from typing import Union
@@ -7,6 +8,7 @@
 
 from gribmagic.unity.core import run_model_download
 from gribmagic.unity.enumerations.weather_models import WeatherModels
+from gribmagic.unity.modules.config.parse_configurations import parse_model_config
 from gribmagic.util import setup_logging
 
 logger = logging.getLogger(__name__)
@@ -20,14 +22,27 @@
 
 @click.group()
 @click.pass_context
-def main(ctx):
+def cli(ctx):
     pass
 
 
-@main.command()
+@cli.group(help="Unified NWP data downloader")
+@click.pass_context
+def unity(ctx):
+    pass
+
+
+@unity.command(name="list", help="List available NWP data labels")
+def unity_list():
+    model_config = parse_model_config()
+    labels = [label for label in model_config.keys() if not label.endswith("-base")]
+    print(json.dumps(labels, indent=4))
+
+
+@unity.command(name="acquire", help="Acquire NWP data")
 @click.option("--model", required=True, help="The weather model name.")
 @click.option("--timestamp", required=True, help="The initialization timestamp.")
-def unity(model: Union[str, WeatherModels], timestamp: Union[str, datetime]):
+def unity_acquire(model: Union[str, WeatherModels], timestamp: Union[str, datetime]):
     logger.info("Starting GribMagic")
     results = run_model_download(
         weather_model=model, initialization_timestamp=timestamp
diff --git a/setup.py b/setup.py
index 4de7aab..8048ab7 100644
--- a/setup.py
+++ b/setup.py
@@ -73,7 +73,7 @@
       },
       entry_points={
         'console_scripts': [
-          'gribmagic = gribmagic.commands:main',
+          'gribmagic = gribmagic.commands:cli',
         ],
       },
 )
diff --git a/tests/unity/test_command.py b/tests/unity/test_commands.py
similarity index 78%
rename from tests/unity/test_command.py
rename to tests/unity/test_commands.py
index c7d28a7..8eb3c65 100644
--- a/tests/unity/test_command.py
+++ b/tests/unity/test_commands.py
@@ -1,3 +1,4 @@
+import json
 import logging
 import os
 import re
@@ -10,7 +11,7 @@
 import responses
 from click.testing import CliRunner
 
-from gribmagic.commands import main
+from gribmagic.commands import cli
 from gribmagic.unity.enumerations.weather_models import WeatherModels
 from gribmagic.unity.models import MODEL_CONFIG
 from gribmagic.unity.modules.config.constants import KEY_VARIABLES
@@ -19,12 +20,37 @@
 model_config.update({KEY_VARIABLES: ["air_temperature_2m", "relative_humidity_2m"]})
 
 
+def test_command_gribmagic_unity_list(capsys):
+    runner = CliRunner()
+    result = runner.invoke(
+        cli=cli,
+        args=[
+            "unity",
+            "list",
+        ],
+        # catch_exceptions=False,
+    )
+    labels = json.loads(result.output)
+    assert labels == [
+        "ncep-gfs-025",
+        "ncep-gfs-050",
+        "ncep-gfs-100",
+        "dwd-icon-global",
+        "dwd-icon-eu",
+        "dwd-icon-eu-eps",
+        "dwd-cosmo-d2",
+        "dwd-cosmo-d2-eps",
+        "meteo-france-arome",
+        "knmi-harmonie",
+    ]
+
+
 @responses.activate
 @patch(
     "gribmagic.unity.models.MODEL_CONFIG",
     {WeatherModels.DWD_ICON_EU.value: model_config},
 )
-def test_command_gribmagic_unity(caplog):
+def test_command_gribmagic_unity_acquire(caplog):
 
     responses.add(
         method=responses.GET,
@@ -40,9 +66,10 @@ def test_command_gribmagic_unity(caplog):
         with caplog.at_level(logging.DEBUG):
             runner = CliRunner()
             result = runner.invoke(
-                cli=main,
+                cli=cli,
                 args=[
                     "unity",
+                    "acquire",
                     "--model=dwd-icon-eu",
                     "--timestamp=2021-10-03T00:00:00Z",
                 ],
@@ -52,7 +79,6 @@ def test_command_gribmagic_unity(caplog):
     # result.exit_code == 1 -- why!?
     # assert result.exit_code == 0
 
-    # assert result.output == 'Hello Peter!\n'
     assert "Starting GribMagic" in caplog.text
     assert (
         "WeatherModels.DWD_ICON_EU: Accessing parameter 'air_temperature_2m'"