Source code for ewoksmx.tests.shell_utils.test_execute_slurm

import logging
import shutil
import subprocess
import tempfile

import pytest

from ewoksmx.shell_utils.execute_slurm import execute_bash_commands




def sbatch_available() -> bool:
    """Check if the sbatch command is available."""
    return subprocess.call(["which", "sbatch"]) == 0





@pytest.fixture
def slurm_parameters():
    if not sbatch_available():
        pytest.skip(reason="Slurm sbatch not available")

    parameters = {
        "time": "00:05:00",
        "ntasks": 1,
        "cpus-per-task": 1,
        "mem": "1GB",
        # "partition": "nice",
    }
    working_directory = tempfile.mkdtemp(prefix="slurm_test_", dir="/tmp_14_days")
    return {"working_directory": working_directory, "parameters": parameters}





def test_execute_bash_commands_success(caplog, capsys, slurm_parameters):
    commands = [
        'echo "Job $SLURM_JOBID start!"',
        'echo "This is a test message for STDOUT"',
        '>&2 echo "This is a test message for STDERR"',
        "hostname",
        'echo "Job finished!"',
    ]

    result = execute_bash_commands(commands, **slurm_parameters)

    result.print()

    with caplog.at_level(logging.DEBUG):
        result.log()

    result.raise_on_error()

    assert "This is a test message for STDOUT" in (result.read_log() or "")
    assert "This is a test message for STDERR" in (result.read_log() or "")
    assert " start!" in (result.read_log() or "")
    assert "Job finished!" in (result.read_log() or "")

    captured = capsys.readouterr()
    assert "This is a test message for STDOUT" in captured.out
    assert "This is a test message for STDERR" in captured.out

    assert any(
        "This is a test message for STDOUT" in record.message
        for record in caplog.records
    )
    assert any(
        "This is a test message for STDERR" in record.message
        for record in caplog.records
    )
    shutil.rmtree(slurm_parameters["working_directory"])





def test_execute_bash_commands_failure(slurm_parameters):
    commands = [
        'echo "Job $SLURM_JOBID start!"',
        'echo "This is a test message for STDOUT"',
        '>&2 echo "This is a test message for STDERR"',
        "ls /cause_error",
        'echo "Job finished!"',
    ]

    result = execute_bash_commands(commands, **slurm_parameters)

    with pytest.raises(RuntimeError):
        result.raise_on_error()
    shutil.rmtree(slurm_parameters["working_directory"])