improve test_slurm_submit(), add tests/test_init.py

Author: janosh

tests/test_init.py

@@ -0,0 +1,36 @@
+import os
+from typing import Any
+
+from mb_discovery import as_dict_handler, chunks
+
+
+def test_has_root_pkg_dir() -> None:
+
+    from mb_discovery import PKG_DIR, ROOT
+
+    assert os.path.isdir(ROOT)
+    assert os.path.isdir(PKG_DIR)
+
+
+def test_chunks() -> None:
+
+    assert list(chunks([], 1)) == []
+    assert list(chunks([1], 1)) == [[1]]
+    assert list(chunks([1, 2], 1)) == [[1], [2]]
+    assert list(chunks([1, 2, 3], 1)) == [[1], [2], [3]]
+    assert list(chunks([1, 2, 3], 2)) == [[1, 2], [3]]
+    assert list(chunks(range(1, 4), 2)) == [range(1, 3), range(3, 4)]
+    assert list(chunks(range(1, 5), 2)) == [range(1, 3), range(3, 5)]
+    assert list(chunks(range(1, 5), 3)) == [range(1, 4), range(4, 5)]
+
+
+def test_as_dict_handler() -> None:
+    class C:
+        def as_dict(self) -> dict[str, Any]:
+            return {"foo": "bar"}
+
+    assert as_dict_handler(C()) == {"foo": "bar"}
+    assert as_dict_handler(1) is None
+    assert as_dict_handler("foo") is None
+    assert as_dict_handler([1, 2, 3]) is None
+    assert as_dict_handler({"foo": "bar"}) is None

tests/test_slurm.py

@@ -2,24 +2,42 @@
 from unittest.mock import patch
 
 import pytest
+from pytest import CaptureFixture
 
 from mb_discovery.slurm import _get_calling_file_path, slurm_submit_python
 
 
-def test_slurm_submit() -> None:
+def test_slurm_submit(capsys: CaptureFixture[str]) -> None:
+    # check slurm_submit_python() does nothing in normal mode
 
+    kwargs = dict(
+        job_name="test_job",
+        log_dir="test_log_dir",
+        time="0:0:1",
+        slurm_flags=("--test-flag",),
+    )
+    slurm_submit_python(**kwargs)  # type: ignore
+
+    stdout, stderr = capsys.readouterr()
+    assert stderr == stderr == ""
+
+    # check slurm_submit_python() does prints cmd calls subprocess.run() in submit mode
     sys.argv += ["slurm-submit"]
-    with pytest.raises(SystemExit) as exc_info, patch(
+
+    with pytest.raises(SystemExit), patch(
         "mb_discovery.slurm.subprocess.run"
-    ) as mock_run:
-        slurm_submit_python(
-            job_name="test_job",
-            log_dir="test_log_dir",
-            time="0:0:1",
-            slurm_flags=("--test_flag",),
-        )
-        assert exc_info.value.code == 0
-        assert mock_run.call_count == 1
+    ) as mock_subprocess_run:
+        slurm_submit_python(**kwargs)  # type: ignore
+
+    assert mock_subprocess_run.call_count == 1
+
+    sbatch_cmd = (
+        "sbatch --partition=icelake --account=LEE-SL3-CPU --time=0:0:1 --job-name "
+        "test_job --output test_log_dir/slurm-%A-%a.out --test-flag --wrap python"
+    )
+    stdout, stderr = capsys.readouterr()
+    assert sbatch_cmd in stdout
+    assert stderr == ""
 
 
 def test_get_calling_file_path() -> None: