From f582cabe198a4ae9db9294b9e47fea887280606a Mon Sep 17 00:00:00 2001
From: cyliang368 <cyliang368@gmail.com>
Date: Sun, 16 Jun 2024 12:34:09 +0200
Subject: [PATCH] Update cli and its test

---
 python/grass/benchmark/app.py                 | 25 ++++++++++++--
 .../benchmark/testsuite/test_benchmark_cli.py | 33 +++++++++++++------
 2 files changed, 46 insertions(+), 12 deletions(-)

diff --git a/python/grass/benchmark/app.py b/python/grass/benchmark/app.py
index 4853e17a4dd..09cdf3f9481 100644
--- a/python/grass/benchmark/app.py
+++ b/python/grass/benchmark/app.py
@@ -69,6 +69,7 @@ def plot_nprocs_cli(args):
         results.results,
         filename=args.output,
         title=args.title,
+        metric=args.metric,
     )
 
 
@@ -161,9 +162,17 @@ def add_results_subcommand(parent_subparsers):
 
 def add_plot_io_arguments(parser):
     """Add input and output arguments to *parser*."""
-    parser.add_argument("input", help="file with results (JSON)", metavar="input_file")
     parser.add_argument(
-        "output", help="output file (e.g., PNG)", nargs="?", metavar="output_file"
+        "input", help="file with results (e.g. results.json)", metavar="input_file"
+    )
+    parser.add_argument(
+        "output",
+        help=(
+            "output file with extension (e.g., figure.png)."
+            " If not provided, the plot will be opened in a new window."
+        ),
+        nargs="?",
+        metavar="output_file",
     )
 
 
@@ -176,6 +185,16 @@ def add_plot_title_argument(parser):
     )
 
 
+def add_plot_metric_argument(parser):
+    """Add metric argument to *parser*."""
+    parser.add_argument(
+        "--metric",
+        help="Metric for the plot (default: time)",
+        default="time",
+        choices=["time", "speedup", "efficiency"],
+    )
+
+
 def add_plot_subcommand(parent_subparsers):
     """Add plot subcommand."""
     main_parser = add_subcommand_parser(
@@ -198,6 +217,7 @@ def add_plot_subcommand(parent_subparsers):
     )
     add_plot_io_arguments(nprocs)
     add_plot_title_argument(nprocs)
+    add_plot_metric_argument(nprocs)
     nprocs.set_defaults(handler=plot_nprocs_cli)
 
 
@@ -206,6 +226,7 @@ def define_arguments():
     parser = argparse.ArgumentParser(
         description="Process results from module benchmarks.",
         prog=get_executable_name(),
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
     )
     subparsers = add_subparsers(parser, dest="command")
 
diff --git a/python/grass/benchmark/testsuite/test_benchmark_cli.py b/python/grass/benchmark/testsuite/test_benchmark_cli.py
index 25765c6a701..8e963245250 100644
--- a/python/grass/benchmark/testsuite/test_benchmark_cli.py
+++ b/python/grass/benchmark/testsuite/test_benchmark_cli.py
@@ -43,14 +43,14 @@ class TestBenchmarkCLI(TestCase):
     """Tests that benchmarkin CLI works"""
 
     json_filename = "plot_test.json"
-    png_filename1 = "plot_test1.png"
-    png_filename2 = "plot_test2.png"
+    png_filenames = [f"plot_test1_{i}.png" for i in range(4)]
+    png_filenames.append("plot_test2.png")
 
     def tearDown(self):
         """Remove test files"""
         remove_file(self.json_filename)
-        remove_file(self.png_filename1)
-        remove_file(self.png_filename2)
+        for filename in self.png_filenames:
+            remove_file(filename)
 
     def test_plot_nprocs_workflow(self):
         """Test that plot nprocs workflow runs"""
@@ -67,8 +67,15 @@ def test_plot_nprocs_workflow(self):
         except grass.exceptions.ParameterError:
             self.skipTest("r.univar without nprocs parameter")
         save_results_to_file([result], self.json_filename)
-        benchmark_main(["plot", "nprocs", self.json_filename, self.png_filename1])
-        self.assertTrue(Path(self.png_filename1).is_file())
+
+        metrics = ["time", "speedup", "efficiency"]
+        benchmark_main(["plot", "nprocs", self.json_filename, self.png_filenames[0]])
+        for png_fname, metric in zip(self.png_filenames[1:4], metrics):
+            benchmark_main(
+                ["plot", "nprocs", "--metric", metric, self.json_filename, png_fname]
+            )
+        for filename in self.png_filenames[:4]:
+            self.assertTrue(Path(filename).is_file())
 
     def test_plot_cells_workflow(self):
         """Test that plot cells workflow runs"""
@@ -82,12 +89,18 @@ def test_plot_cells_workflow(self):
             resolutions=[1000, 500],
         )
         save_results_to_file([result], self.json_filename)
-        benchmark_main(["plot", "cells", self.json_filename, self.png_filename1])
-        self.assertTrue(Path(self.png_filename1).is_file())
+        benchmark_main(["plot", "cells", self.json_filename, self.png_filenames[0]])
+        self.assertTrue(Path(self.png_filenames[0]).is_file())
         benchmark_main(
-            ["plot", "cells", "--resolutions", self.json_filename, self.png_filename2]
+            [
+                "plot",
+                "cells",
+                "--resolutions",
+                self.json_filename,
+                self.png_filenames[-1],
+            ]
         )
-        self.assertTrue(Path(self.png_filename2).is_file())
+        self.assertTrue(Path(self.png_filenames[-1]).is_file())
 
 
 if __name__ == "__main__":