From 2f1562801c0297d07f3ce2ed5ecd84a5fa2144fd Mon Sep 17 00:00:00 2001
From: Nitin Garg <gargnitin@google.com>
Date: Fri, 2 Aug 2024 12:13:42 +0000
Subject: [PATCH] Fix argparser and dry-run support in run_tests scripts

---
 .../testing_on_gke/examples/dlio/run_tests.py | 21 ++++++++++------
 .../testing_on_gke/examples/fio/run_tests.py  | 25 ++++++++++---------
 2 files changed, 26 insertions(+), 20 deletions(-)

diff --git a/perfmetrics/scripts/testing_on_gke/examples/dlio/run_tests.py b/perfmetrics/scripts/testing_on_gke/examples/dlio/run_tests.py
index 4b59cf42f6..a2a2354e93 100644
--- a/perfmetrics/scripts/testing_on_gke/examples/dlio/run_tests.py
+++ b/perfmetrics/scripts/testing_on_gke/examples/dlio/run_tests.py
@@ -18,11 +18,7 @@
 """This program takes in a json dlio test-config file and generates and deploys helm charts."""
 
 import argparse
-from collections.abc import Sequence
-import os
 import subprocess
-
-from absl import app
 import dlio_workload
 
 
@@ -60,18 +56,27 @@ def main(args) -> None:
   helmInstallCommands = createHelmInstallCommands(dlioWorkloads)
   for helmInstallCommand in helmInstallCommands:
     print(f'{helmInstallCommand}')
-    run_command(helmInstallCommand)
+    if not args.dry_run:
+      run_command(helmInstallCommand)
 
 
 if __name__ == '__main__':
   parser = argparse.ArgumentParser(
-      prog='DLIO test runner',
+      prog='DLIO Unet3d test runner',
       description=(
           'This program takes in a json dlio test-config file and generates'
-          ' helm install commands.'
+          ' helm install commands to execute them using the active GKE cluster.'
       ),
-      # epilog='Text at the bottom of help',
   )
   parser.add_argument('--workload-config')  # positional argument
+  parser.add_argument(
+      '-n',
+      '--dry-run',
+      action='store_true',
+      help=(
+          'Only print out the test configurations that will run,'
+          ' not actually run them.'
+      ),
+  )
   args = parser.parse_args()
   main(args)
diff --git a/perfmetrics/scripts/testing_on_gke/examples/fio/run_tests.py b/perfmetrics/scripts/testing_on_gke/examples/fio/run_tests.py
index 32d374d218..920526a050 100644
--- a/perfmetrics/scripts/testing_on_gke/examples/fio/run_tests.py
+++ b/perfmetrics/scripts/testing_on_gke/examples/fio/run_tests.py
@@ -18,13 +18,7 @@
 """This program takes in a json fio test-config file and generates and deploys helm charts."""
 
 import argparse
-from collections.abc import Sequence
-import json
-import os
-import pprint
 import subprocess
-
-from absl import app
 import fio_workload
 
 
@@ -61,23 +55,30 @@ def main(args) -> None:
   fioWorkloads = fio_workload.ParseTestConfigForFioWorkloads(
       args.workload_config
   )
-  # for fioWorkload in fioWorkloads:
-  # fioWorkload.PPrint()
   helmInstallCommands = createHelmInstallCommands(fioWorkloads)
   for helmInstallCommand in helmInstallCommands:
     print(f'{helmInstallCommand}')
-    run_command(f'{helmInstallCommand}')
+    if not args.dry_run:
+      run_command(helmInstallCommand)
 
 
 if __name__ == '__main__':
   parser = argparse.ArgumentParser(
       prog='FIO test runner',
       description=(
-          'This program takes in a json fio test-config file and generates'
-          ' helm install commands.'
+          'This program takes in a json test-config file and generates'
+          ' helm install commands to execute them using the active GKE cluster.'
       ),
-      # epilog='Text at the bottom of help',
   )
   parser.add_argument('--workload-config')
+  parser.add_argument(
+      '-n',
+      '--dry-run',
+      action='store_true',
+      help=(
+          'Only print out the test configurations that will run,'
+          ' not actually run them.'
+      ),
+  )
   args = parser.parse_args()
   main(args)