Add test of update argument expansion, fix incorrect types found via

the test
LLNL · Nov 22, 2024 · 63220bb · 63220bb
1 parent b3d539a
commit 63220bb
Show file tree

Hide file tree

Showing 2 changed files with 56 additions and 4 deletions.
diff --git a/maestrowf/maestro.py b/maestrowf/maestro.py
@@ -559,23 +559,23 @@ def setup_argparser():
     update.add_argument(
         "--throttle",
         action="append",
-        type=str,
+        type=int,
         default=[],
         help="Update the maximum number of inflight jobs allowed to execute"
         " simultaneously (0 denotes no throttling)."
     )
     update.add_argument(
         "--rlimit",
         action="append",
-        type=str,
+        type=int,
         default=[],
         help="Update the maximum number of restarts allowed when steps "
         "specify a restart command (0 denotes no limit)."
     )
     update.add_argument(
         "--sleep",
         action="append",
-        type=str,
+        type=int,
         default=[],
         help="Update the time (in seconds) that the manager will "
         "wait between job status checks.",

diff --git a/tests/test_cli.py b/tests/test_cli.py
@@ -13,7 +13,7 @@
         (["--rlimit", "2", "--sleep", "30", 'study_workspace_1', 'study_workspace_2'], True),
         (["--rlimit", "2", "--sleep", "30", "--sleep", "32", 'study_workspace_1', 'study_workspace_2'], True),
         (["--rlimit", "2", "--sleep", "30", "--sleep", "32", "--sleep", "33", 'study_workspace_1', 'study_workspace_2'], False),
-        (["--rlimit", "2", "--sleep", "30", "--sleep", "32", 'study_workspace_1', 'study_workspace_2', 'study_workspace_3'], True),
+        (["--rlimit", "2", "--sleep", "30", "--sleep", "32", 'study_workspace_1', 'study_workspace_2', 'study_workspace_3'], False),
     ]
 )
 def test_validate_update_args(cli_args, args_are_valid):
@@ -28,3 +28,55 @@ def test_validate_update_args(cli_args, args_are_valid):
     args = parser.parse_args(maestro_cli)
     # assert args is None
     assert maestro.validate_update_args(args, args.directory) is args_are_valid
+
+
+@pytest.mark.parametrize(
+    "cli_args, expected_expanded_args",
+    [
+        (
+            ["--rlimit", "2", 'study_workspace_1'],
+            [{'rlimit': 2, 'throttle': None, 'sleep': None}, ]
+        ),
+        (
+            ["--rlimit", "2", 'study_workspace_1', 'study_workspace_2'],
+            [
+                {'rlimit': 2, 'throttle': None, 'sleep': None},
+                {'rlimit': 2, 'throttle': None, 'sleep': None}
+            ]
+        ),
+        (
+            ["--rlimit", "2", "--sleep", "30", 'study_workspace_1'],
+            [{'rlimit': 2, 'sleep': 30, 'throttle': None}, ]
+        ),
+        (
+            ["--rlimit", "2", "--sleep", "30", 'study_workspace_1', 'study_workspace_2'],
+            [
+                {'rlimit': 2, 'sleep': 30, 'throttle': None},
+                {'rlimit': 2, 'sleep': 30, 'throttle': None}
+            ]
+        ),
+        (
+            ["--rlimit", "2", "--sleep", "30", "--sleep", "32", 'study_workspace_1', 'study_workspace_2'],
+            [
+                {'rlimit': 2, 'sleep': 30, 'throttle': None},
+                {'rlimit': 2, 'sleep': 32, 'throttle': None}
+            ],
+        ),
+    ]
+)
+def test_expand_update_args(cli_args, expected_expanded_args):
+    """
+    Test expansion of arguments passed to the 'maestro update' cli command,
+    i.e. repeating single values for all study workspaces being updated
+    """
+    parser = maestro.setup_argparser()
+    maestro_cli = ["update"]
+
+    maestro_cli.extend(cli_args)
+    print(f"{maestro_cli=}")
+    args = parser.parse_args(maestro_cli)
+    expanded_args = maestro.expand_update_args(args, args.directory)
+    print(f"{args=}")
+    print(f"{expanded_args=}")
+    print(f"{expected_expanded_args=}")
+    assert expanded_args == expected_expanded_args