Merge pull request #13 from larsbuntemeyer/parser-tests

Parser tests

Author: larsbuntemeyer

enjoy_slurm/slurm.py

@@ -2,6 +2,7 @@
 
 from .utils import (
     kwargs_to_list,
+    args_to_list,
     parse_sacct,
     execute,
     create_scontrol_func,
@@ -31,13 +32,20 @@ def sbatch(jobscript=None, *args, **kwargs):
     jobid : int
         Slurm jobid.
 
+
     """
     if jobscript is None:
         jobscript = []
     else:
         jobscript = [jobscript]
 
-    command = ["sbatch", "--parsable"] + list(args) + kwargs_to_list(kwargs) + jobscript
+    command = (
+        ["sbatch", "--parsable"]
+        + args_to_list(args)
+        + kwargs_to_list(kwargs)
+        + jobscript
+    )
+
     jobid = int(execute(command))
 
     return jobid

enjoy_slurm/utils.py

@@ -43,9 +43,26 @@ def parse_slurm_arg(a):
     """parse slurm arguments to str or list of str with colons"""
     if isinstance(a, list):
         return ":".join([str(x) for x in a])
+    if a is True:
+        return []
     return [str(a)]
 
 
+def args_to_list(args):
+    """parse sbatch arguments to list"""
+    args_list = []
+    for a in args:
+        if isinstance(a, list):
+            args_list += a
+        elif "=" in a:
+            args_list += a.split("=")
+        elif " " in a:
+            args_list += a.split()
+        else:
+            args_list.append(a)
+    return args_list
+
+
 def kwargs_to_list(d):
     """parse arguments to command line arguments for sbatch"""
     r = []
@@ -54,12 +71,21 @@ def kwargs_to_list(d):
         if k in own_kwargs:
             continue
         flag = "--" + k.replace("_", "-")
-        r += [flag]
+
         if k == "dependency" and v is not None:
-            # flag += "=" + parse_dependency(v, d.get("how", None))
+            r += [flag]
             r += parse_dependency(v, d.get("how", None))
-        elif v is not True:
+            continue
+        if k == "kill_on_invalid_dep" and not isinstance(v, str):
+            r += [flag]
+            r += ["no"] if v is False else ["yes"]
+            continue
+        if v:
+            r += [flag]
             r += parse_slurm_arg(v)
+            continue
+        if v is False:
+            continue
         # r += [flag]
     return r
 

tests/test_parser.py

@@ -1,6 +1,11 @@
 import pytest
 
-from enjoy_slurm.utils import parse_dependency, kwargs_to_list, handle_sacct_format
+from enjoy_slurm.utils import (
+    parse_dependency,
+    kwargs_to_list,
+    handle_sacct_format,
+    args_to_list,
+)
 from enjoy_slurm.config import default_sacct_format
 
 
@@ -12,6 +17,10 @@ def test_parse_dependency():
 
 
 def test_kwargs_to_list():
+    assert kwargs_to_list({"hello": True}) == ["--hello"]
+    assert kwargs_to_list({"hello": False}) == []
+    assert kwargs_to_list({"hello": 1}) == ["--hello", "1"]
+
     kwargs = {"partition": "test", "dependency": [1, 2, 3]}
     assert kwargs_to_list(kwargs) == [
         "--partition",
@@ -26,6 +35,41 @@ def test_kwargs_to_list():
         "--dependency",
         "afterany:1:2:3",
     ]
+    kwargs = {
+        "partition": "test",
+        "dependency": [1, 2, 3],
+        "how": "afterany",
+        "kill_on_invalid_dep": True,
+    }
+    assert kwargs_to_list(kwargs) == [
+        "--partition",
+        "test",
+        "--dependency",
+        "afterany:1:2:3",
+        "--kill-on-invalid-dep",
+        "yes",
+    ]
+    kwargs = {"kill_on_invalid_dep": False}
+    assert kwargs_to_list(kwargs) == ["--kill-on-invalid-dep", "no"]
+    kwargs = {"kill_on_invalid_dep": "yes"}
+    assert kwargs_to_list(kwargs) == ["--kill-on-invalid-dep", "yes"]
+
+
+def test_args_to_list():
+    assert args_to_list(("--partition=shared",)) == ["--partition", "shared"]
+    assert args_to_list(("--partition shared",)) == ["--partition", "shared"]
+    assert args_to_list(("--partition shared", "--dependency=afterok:1:2:3")) == [
+        "--partition",
+        "shared",
+        "--dependency",
+        "afterok:1:2:3",
+    ]
+    assert args_to_list(("--partition shared", ["--dependency", "afterok:1:2:3"])) == [
+        "--partition",
+        "shared",
+        "--dependency",
+        "afterok:1:2:3",
+    ]
 
 
 def test_sacct_format():

tests/test_with_slurm.py

@@ -10,6 +10,11 @@
 
 print(hostname)
 
+if on_levante:
+    kwargs = {"partition": "shared", "account": "ch0636"}
+else:
+    kwargs = {}
+
 
 @requires_slurm
 def test_sbatch():
@@ -30,3 +35,11 @@ def test_partitions():
     for p in partitions:
         pdict = slurm.scontrol.show(partition=p)
         assert p in pdict
+
+
+@requires_slurm
+def test_sacct():
+    jobid = slurm.sbatch(wrap="echo Hello World", **kwargs)
+    acct = slurm.sacct(jobid=jobid)
+    scon = slurm.scontrol.show(jobid=jobid)
+    assert str(jobid) in scon