diff --git a/junifer/api/functions.py b/junifer/api/functions.py index 90ea8f839..21d2aa78a 100644 --- a/junifer/api/functions.py +++ b/junifer/api/functions.py @@ -84,7 +84,7 @@ def run( # Convert str to Path if isinstance(workdir, str): workdir = Path(workdir) - if not isinstance(elements, List) and elements is not None: + if not isinstance(elements, list) and elements is not None: elements = [elements] # Get datagrabber to use datagrabber_object = _get_datagrabber(datagrabber) @@ -186,7 +186,7 @@ def queue( jobdir = cwd / "junifer_jobs" / jobname logger.info(f"Creating job in {str(jobdir.absolute())}") if jobdir.exists(): - if overwrite is not True: + if not overwrite: raise_error( f"Job folder for {jobname} already exists. " "This error is raised to prevent overwriting job files " @@ -234,7 +234,7 @@ def queue( elements = dg.get_elements() # TODO: Fix typing of elements - if not isinstance(elements, List): + if not isinstance(elements, list): elements = [elements] # type: ignore typing.cast(List[Union[str, Tuple]], elements) @@ -258,7 +258,7 @@ def queue( **kwargs, ) else: - raise ValueError(f"Unknown queue kind: {kind}") + raise_error(f"Unknown queue kind: {kind}") logger.info("Queue done") diff --git a/junifer/api/tests/test_functions.py b/junifer/api/tests/test_functions.py index c69c939b0..c1ffd3d5d 100644 --- a/junifer/api/tests/test_functions.py +++ b/junifer/api/tests/test_functions.py @@ -5,15 +5,19 @@ # Synchon Mandal # License: AGPL +import logging from pathlib import Path +from typing import List, Tuple, Union import pytest +import yaml import junifer.testing.registry # noqa: F401 -from junifer.api.functions import collect, run +from junifer.api.functions import collect, queue, run from junifer.datagrabber.base import BaseDataGrabber from junifer.pipeline.registry import build + # Define datagrabber datagrabber = { "kind": "OasisVBMTestingDatagrabber", @@ -144,10 +148,562 @@ def test_run_and_collect(tmp_path: Path) -> None: assert uri.exists() -@pytest.mark.skip(reason="HTCondor not installed on system.") -def test_queue_condor() -> None: - """Test job queueing in HTCondor.""" - pass +def test_queue_correct_yaml_config( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Test proper YAML config generation for queueing. + + Parameters + ---------- + tmp_path : pathlib.Path + The path to the test directory. + monkeypatch : pytest.MonkeyPatch + The monkeypatch object. + + """ + with monkeypatch.context() as m: + m.chdir(tmp_path) + queue( + config={ + "with": "junifer.testing.registry", + "workdir": str(Path(tmp_path).resolve()), + "datagrabber": datagrabber, + "markers": markers, + "storage": storage, + "env": { + "kind": "conda", + "name": "junifer", + }, + "mem": "8G", + }, + kind="HTCondor", + jobname="yaml_config_gen_check", + ) + + generated_config_yaml_path = Path( + tmp_path / "junifer_jobs" / "yaml_config_gen_check" / "config.yaml" + ) + with open(generated_config_yaml_path, "r") as f: + yaml_config = yaml.unsafe_load(f) + # Check for correct YAML config generation + assert all( + key in yaml_config.keys() + for key in [ + "with", + "workdir", + "datagrabber", + "markers", + "storage", + "env", + "mem", + ] + ) + assert "queue" not in yaml_config.keys() + + +def test_queue_invalid_job_queue( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Test queue function for invalid job queue. + + Parameters + ---------- + tmp_path : pathlib.Path + The path to the test directory. + monkeypatch : pytest.MonkeyPatch + The monkeypatch object. + + """ + with pytest.raises(ValueError, match="Unknown queue kind"): + with monkeypatch.context() as m: + m.chdir(tmp_path) + queue( + config={"elements": ["sub-001"]}, + kind="ABC", + ) + + +def test_queue_assets_disallow_overwrite( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Test overwrite prevention of queue files. + + Parameters + ---------- + tmp_path : pathlib.Path + The path to the test directory. + monkeypatch : pytest.MonkeyPatch + The monkeypatch object. + + """ + with pytest.raises(ValueError, match="Either delete the directory"): + with monkeypatch.context() as m: + m.chdir(tmp_path) + # First generate assets + queue( + config={"elements": ["sub-001"]}, + kind="HTCondor", + jobname="prevent_overwrite", + ) + # Re-run to trigger error + queue( + config={"elements": ["sub-001"]}, + kind="HTCondor", + jobname="prevent_overwrite", + ) + + +def test_queue_assets_allow_overwrite( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + """Test overwriting of queue files. + + Parameters + ---------- + tmp_path : pathlib.Path + The path to the test directory. + monkeypatch : pytest.MonkeyPatch + The monkeypatch object. + caplog : pytest.LogCaptureFixture + The logcapturefixture object. + + """ + with monkeypatch.context() as m: + m.chdir(tmp_path) + # First generate assets + queue( + config={"elements": ["sub-001"]}, + kind="HTCondor", + jobname="allow_overwrite", + ) + with caplog.at_level(logging.INFO): + # Re-run to overwrite + queue( + config={"elements": ["sub-001"]}, + kind="HTCondor", + jobname="allow_overwrite", + overwrite=True, + ) + assert "Deleting existing job directory" in caplog.text + + +@pytest.mark.parametrize( + "with_", + [ + "a.py", + ["a.py"], + ["a.py", "b"], + ], +) +def test_queue_with_imports( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + with_: Union[str, List[str]], +) -> None: + """Test queue with `with` imports. + + Parameters + ---------- + tmp_path : pathlib.Path + The path to the test directory. + monkeypatch : pytest.MonkeyPatch + The monkeypatch object. + caplog : pytest.LogCaptureFixture + The logcapturefixture object. + with_ : str or list of str + The parametrized imports. + + """ + with monkeypatch.context() as m: + m.chdir(tmp_path) + # Create test file, keeping it simple without conditionals + (tmp_path / "a.py").touch() + with caplog.at_level(logging.DEBUG): + queue( + config={"with": with_}, + kind="HTCondor", + jobname="with_import_check", + elements="sub-001", + ) + assert "Copying" in caplog.text + assert "Queue done" in caplog.text + + # Check that file is copied + assert Path( + tmp_path / "junifer_jobs" / "with_import_check" / "a.py" + ).is_file() + + +@pytest.mark.parametrize( + "elements", + [ + "sub-001", + ["sub-001"], + ["sub-001", "sub-002"], + ("sub-001", "ses-001"), + [("sub-001", "ses-001")], + [("sub-001", "ses-001"), ("sub-001", "ses-002")], + ], +) +def test_queue_with_elements( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + elements: Union[str, List[Union[str, Tuple[str]]], Tuple[str]], +) -> None: + """Test queue with elements. + + Parameters + ---------- + tmp_path : pathlib.Path + The path to the test directory. + monkeypatch : pytest.MonkeyPatch + The monkeypatch object. + caplog : pytest.LogCaptureFixture + The logcapturefixture object. + elements : str of list of str + The parametrized elements for the queue. + + """ + with monkeypatch.context() as m: + m.chdir(tmp_path) + with caplog.at_level(logging.INFO): + queue( + config={}, + kind="HTCondor", + elements=elements, + ) + assert "Queue done" in caplog.text + + +def test_queue_without_elements( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + """Test queue without elements. + + Parameters + ---------- + tmp_path : pathlib.Path + The path to the test directory. + monkeypatch : pytest.MonkeyPatch + The monkeypatch object. + caplog : pytest.LogCaptureFixture + The logcapturefixture object. + + """ + with monkeypatch.context() as m: + m.chdir(tmp_path) + with caplog.at_level(logging.INFO): + queue( + config={"datagrabber": datagrabber}, + kind="SLURM", + ) + assert "Queue done" in caplog.text + + +def test_queue_condor_invalid_python_env( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Test invalid Python environment check for HTCondor. + + Parameters + ---------- + tmp_path : pathlib.Path + The path to the test directory. + monkeypatch : pytest.MonkeyPatch + The monkeypatch object. + + """ + with pytest.raises(ValueError, match="Unknown env kind"): + with monkeypatch.context() as m: + m.chdir(tmp_path) + queue( + config={"elements": "sub-001"}, + kind="HTCondor", + env={"kind": "galaxy"}, + ) + + +def test_queue_condor_conda_python( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + """Test conda Python environment check for HTCondor. + + Parameters + ---------- + tmp_path : pathlib.Path + The path to the test directory. + monkeypatch : pytest.MonkeyPatch + The monkeypatch object. + caplog : pytest.LogCaptureFixture + The logcapturefixture object. + + """ + with monkeypatch.context() as m: + m.chdir(tmp_path) + with caplog.at_level(logging.INFO): + queue( + config={"elements": "sub-001"}, + kind="HTCondor", + jobname="conda_env_check", + env={"kind": "conda", "name": "conda-env"}, + ) + assert "Copying" in caplog.text + assert ( + Path( + tmp_path + / "junifer_jobs" + / "conda_env_check" + / "run_conda.sh" + ) + .stat() + .st_mode + == 33261 + ) + + +def test_queue_condor_venv_python( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + """Test venv Python environment check for HTCondor. + + Parameters + ---------- + tmp_path : pathlib.Path + The path to the test directory. + monkeypatch : pytest.MonkeyPatch + The monkeypatch object. + caplog : pytest.LogCaptureFixture + The logcapturefixture object. + + """ + with monkeypatch.context() as m: + m.chdir(tmp_path) + with caplog.at_level(logging.INFO): + queue( + config={"elements": "sub-001"}, + kind="HTCondor", + env={"kind": "venv", "name": "venv-env"}, + ) + # TODO: needs implementation for testing + + +@pytest.mark.parametrize( + "elements, env, mem, cpus, disk, collect", + [ + ( + ["sub-001"], + {"kind": "conda", "name": "conda-env"}, + "4G", + 4, + "4G", + True, + ), + ( + ["sub-001"], + {"kind": "venv", "name": "venv-env"}, + "8G", + 8, + "8G", + False, + ), + (["sub-001"], {"kind": "local"}, "12G", 12, "12G", True), + ], +) +def test_queue_condor_assets_generation( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + elements: str, + env: str, + mem: str, + cpus: int, + disk: str, + collect: bool, +) -> None: + """Test HTCondor generated assets. + + Parameters + ---------- + tmp_path : pathlib.Path + The path to the test directory. + monkeypatch : pytest.MonkeyPatch + The monkeypatch object. + caplog : pytest.LogCaptureFixture + The logcapturefixture object. + elements : str + The parametrized element names. + env : dict + The parametrized env names. + mem : str + The parametrized memory size. + cpus : int + The parametrized CPUs. + disk : str + The parametrized disk size. + collect : bool + The parametrized collect option. + + """ + jobname = "condor_assets_gen_check" + with monkeypatch.context() as m: + m.chdir(tmp_path) + with caplog.at_level(logging.INFO): + queue( + config={"elements": elements}, + kind="HTCondor", + jobname=jobname, + env=env, + mem=mem, + cpus=cpus, + disk=disk, + collect=collect, + ) + + # Check log directory creation + assert Path(tmp_path / "junifer_jobs" / jobname / "logs").is_dir() + + run_submit_file_path = Path( + tmp_path / "junifer_jobs" / jobname / f"run_{jobname}.submit" + ) + # Check junifer run submit file + assert run_submit_file_path.is_file() + # Read run submit file to check if resources are correct + with open(run_submit_file_path, "r") as f: + for line in f.read().splitlines(): + if "request_cpus" in line: + assert int(line.split("=")[1].strip()) == cpus + if "request_memory" in line: + assert line.split("=")[1].strip() == mem + if "request_disk" in line: + assert line.split("=")[1].strip() == disk + + # Check junifer collect submit file + assert Path( + tmp_path + / "junifer_jobs" + / jobname + / f"collect_{jobname}.submit" + ).is_file() + + dag_file_path = Path( + tmp_path / "junifer_jobs" / jobname / f"{jobname}.dag" + ) + # Check junifer dag file + assert dag_file_path.is_file() + # Read dag file to check if collect job is found + element_count = 0 + has_collect_job = False + with open(dag_file_path, "r") as f: + for line in f.read().splitlines(): + if "JOB" in line: + element_count += 1 + if "collect" in line: + has_collect_job = True + + if collect: + assert len(elements) == element_count - 1 + assert has_collect_job is True + else: + assert len(elements) == element_count + assert has_collect_job is False + + # Check submit log + assert ( + "HTCondor job files created, to submit the job, run" + in caplog.text + ) + + +def test_queue_condor_extra_preamble( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Test HTCondor extra preamble addition. + + Parameters + ---------- + tmp_path : pathlib.Path + The path to the test directory. + monkeypatch : pytest.MonkeyPatch + The monkeypatch object. + + """ + jobname = "condor_extra_preamble_check" + extra_preamble = "FOO = BAR" + with monkeypatch.context() as m: + m.chdir(tmp_path) + queue( + config={"elements": ["sub-001"]}, + kind="HTCondor", + jobname=jobname, + extra_preamble=extra_preamble, + ) + + # Check extra preamble in run submit file + run_submit_file_path = Path( + tmp_path / "junifer_jobs" / jobname / f"run_{jobname}.submit" + ) + with open(run_submit_file_path, "r") as f: + for line in f.read().splitlines(): + if "FOO" in line: + assert line.strip() == extra_preamble + + # Check extra preamble in collect submit file + collect_submit_file_path = Path( + tmp_path / "junifer_jobs" / jobname / f"collect_{jobname}.submit" + ) + with open(collect_submit_file_path, "r") as f: + for line in f.read().splitlines(): + if "FOO" in line: + assert line.strip() == extra_preamble + + +def test_queue_condor_submission_fail( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + """Test HTCondor job submission failure. + + Parameters + ---------- + tmp_path : pathlib.Path + The path to the test directory. + monkeypatch : pytest.MonkeyPatch + The monkeypatch object. + caplog : pytest.LogCaptureFixture + The logcapturefixture object. + + """ + with pytest.raises( + FileNotFoundError, + match="No such file or directory: 'condor_submit_dag'", + ): + with monkeypatch.context() as m: + m.chdir(tmp_path) + with caplog.at_level(logging.INFO): + queue( + config={"elements": ["sub-001"]}, + kind="HTCondor", + jobname="condor_job_submission_fail", + submit=True, + ) + # Check submit log + assert "Submitting HTCondor job" in caplog.text @pytest.mark.skip(reason="SLURM not installed on system.")