Raising error for reserved keywords under function parameter options in get_allocation

smartsim/wlm/slurm.py

@@ -247,18 +247,24 @@ def _get_alloc_cmd(
     if account:
         salloc_args.extend(["-A", str(account)])
 
+    arguments = set(options.keys() if options is not None else {})
+    invalid = {"t", "time", "N", "nodes", "A", "account"}
+
+    if valid := arguments.intersection(invalid):
+        raise ValueError(
+            f"Expecting time, nodes, account as an argument. Also received: {valid}"
+        )
+
     for opt, val in (options or {}).items():
-        if opt not in ["t", "time", "N", "nodes", "A", "account"]:
-            short_arg = bool(len(str(opt)) == 1)
-            prefix = "-" if short_arg else "--"
-            if not val:
-                salloc_args += [prefix + opt]
+        short_arg = bool(len(str(opt)) == 1)
+        prefix = "-" if short_arg else "--"
+        if not val:
+            salloc_args += [prefix + opt]
+        else:
+            if short_arg:
+                salloc_args += [prefix + opt, str(val)]
             else:
-                if short_arg:
-                    salloc_args += [prefix + opt, str(val)]
-                else:
-                    salloc_args += ["=".join((prefix + opt, str(val)))]
-
+                salloc_args += ["=".join((prefix + opt, str(val)))]
     return salloc_args
 
 

tests/full_wlm/test_slurm_allocation.py

@@ -48,11 +48,31 @@ def test_get_release_allocation_w_options(wlmutils):
     """test slurm interface for obtaining allocations"""
     options = {"ntasks-per-node": 1}
     account = wlmutils.get_test_account()
-    alloc = slurm.get_allocation(nodes=1, time="00:05:00", options=options, account=account)
+    alloc = slurm.get_allocation(
+        nodes=1, time="00:05:00", options=options, account=account
+    )
     time.sleep(5)  # give slurm a rest
     slurm.release_allocation(alloc)
 
 
+@pytest.mark.parametrize(
+    "func_parameters,test_parameters",
+    [
+        ("t", "T"),
+        ("time", "TIME"),
+        ("nodes", "NODES"),
+        ("N", "N"),
+        ("account", "ACCOUNT"),
+        ("A", "A"),
+    ],
+)
+def test_get_allocation_bad_params(func_parameters, test_parameters):
+    """test get_allocation with reserved keywords as option"""
+    with pytest.raises(ValueError) as e:
+        alloc = slurm.get_allocation(options={func_parameters: test_parameters})
+    assert "Expecting time, nodes, account as an argument" in e.value.args[0]
+
+
 # --------- Error handling ----------------------------
 
 

tests/test_slurm_get_alloc.py

@@ -47,24 +47,3 @@ def test_get_alloc_format():
         "--ntasks-per-node=5",
     ]
     assert alloc_cmd == result
-
-
-def test_get_alloc_format_overlap():
-    """Test get alloc with collision between arguments and options"""
-    time = "10:00:00"
-    nodes = 5
-    account = "A35311"
-    options = {"N": 10, "time": "15", "account": "S1242"}
-    alloc_cmd = _get_alloc_cmd(nodes, time, account, options)
-    result = [
-        "--no-shell",
-        "-N",
-        "5",
-        "-J",
-        "SmartSim",
-        "-t",
-        "10:00:00",
-        "-A",
-        "A35311",
-    ]
-    assert result == alloc_cmd