test_lib.py

import subprocess
import tempfile
import time
from pathlib import Path
from string import Template

import pytest
from ase import build

import slurm_wrapper


def test_get_slurm_energy_uses_mocked_submission(monkeypatch):
    atoms = build.bulk("Si").repeat((1, 1, 1))

    with tempfile.TemporaryDirectory() as tmpdir:
        working_directory = Path(tmpdir) / "work"

        def fake_submit(slurm_string: str) -> subprocess.CompletedProcess[str]:
            assert str(working_directory) in slurm_string
            assert "#SBATCH --job-name=test" in slurm_string
            time.sleep(0.05)
            (working_directory / "energy.txt").write_text("3.14\n")
            return subprocess.CompletedProcess(
                args=["sbatch", "--parsable", "--wait"],
                returncode=0,
                stdout="12345\n",
                stderr="",
            )

        monkeypatch.setattr(slurm_wrapper, "_run_slurm_submission", fake_submit)

        future = slurm_wrapper.get_slurm_energy(
            atoms,
            Template("#!/bin/bash\n#SBATCH --job-name=test\ncd $template_dir"),
            working_directory,
        )
        energy = future.result()

        assert energy == 3.14
        assert (working_directory / "sbatch.stdout").read_text() == "12345\n"
        assert (working_directory / "sbatch.returncode").read_text() == "0\n"


def test_get_slurm_energy_raises_when_submission_fails(monkeypatch):
    atoms = build.bulk("Si").repeat((1, 1, 1))

    with tempfile.TemporaryDirectory() as tmpdir:
        working_directory = Path(tmpdir) / "work"

        def fake_submit(slurm_string: str) -> subprocess.CompletedProcess[str]:
            assert str(working_directory) in slurm_string
            time.sleep(0.05)
            return subprocess.CompletedProcess(
                args=["sbatch", "--parsable", "--wait"],
                returncode=1,
                stdout="",
                stderr="job failed",
            )

        monkeypatch.setattr(slurm_wrapper, "_run_slurm_submission", fake_submit)

        with pytest.raises(RuntimeError, match="Slurm job failed"):
            future = slurm_wrapper.get_slurm_energy(
                atoms,
                Template("#!/bin/bash\n#SBATCH --job-name=test\ncd $template_dir"),
                working_directory,
            )
            future.result()