diff --git a/src/glicid_spawner/slurm.py b/src/glicid_spawner/slurm.py index 000a0d4..e96eb8e 100644 --- a/src/glicid_spawner/slurm.py +++ b/src/glicid_spawner/slurm.py @@ -6,6 +6,7 @@ import subprocess from dataclasses import dataclass, field from itertools import groupby from operator import attrgetter +from pathlib import Path @dataclass @@ -94,6 +95,12 @@ def sinfo_filter(resources: list, with_states=('idle', 'mixed')) -> dict: return {key: values for key, values in resources.items() if values} +def sinfo_from_file(fname, with_states=('idle', 'mixed')) -> dict: + """SLURM SINFO resources available from a given file.""" + content = Path(fname).read_text() + return sinfo_filter(sinfo_reader(content), with_states=with_states) + + def sinfo(username: str = None, with_states=('idle', 'mixed')) -> dict: """SLURM SINFO resources available for a given user.""" return sinfo_filter(sinfo_reader(sinfo_run(username=username)), with_states=with_states) diff --git a/tests/test_slurm.py b/tests/test_slurm.py index 2c59832..50ebfe0 100644 --- a/tests/test_slurm.py +++ b/tests/test_slurm.py @@ -8,13 +8,15 @@ from glicid_spawner.slurm import ( SlurmNode, sinfo, sinfo_filter, + sinfo_from_file, sinfo_reader, sinfo_run, subprocess, ) DATA = Path(__file__).parent / 'data' -SINFO = (DATA / 'sinfo.txt').read_text() +SINFO_FILE = DATA / 'sinfo.txt' +SINFO_CONTENT = SINFO_FILE.read_text() def test_slurm_dataclasses(): @@ -71,7 +73,7 @@ def test_slurm_sinfo_run(monkeypatch): def test_slurm_sinfo_reader(): """Test SLURM SINFO reader.""" - nodes = sinfo_reader(SINFO) + nodes = sinfo_reader(SINFO_CONTENT) for node in nodes: assert isinstance(node, SlurmNode) @@ -117,7 +119,7 @@ def test_slurm_sinfo_reader(): def test_slurm_sinfo_filter(monkeypatch): """Test SLURM SINFO filtered resources.""" - resources = sinfo_reader(SINFO) + resources = sinfo_reader(SINFO_CONTENT) clusters = sinfo_filter(resources) @@ -156,9 +158,21 @@ def test_slurm_sinfo_filter(monkeypatch): assert [len(partitions) for partitions in clusters.values()] == [1] +def test_slurm_sinfo_from_file(monkeypatch): + """Test SLURM SINFO resources from file.""" + resources = sinfo_from_file(SINFO_FILE, with_states=('idle')) + + assert [ + node.hostname + for cluster, partitions in resources.items() + for nodes in partitions.values() + for node in nodes + ] == ['nazare001', 'gnode2', 'visu1', 'cribbar001'] + + def test_slurm_sinfo_resources(monkeypatch): """Test SLURM SINFO resources.""" - monkeypatch.setattr(subprocess, 'check_output', lambda _: SINFO.encode()) + monkeypatch.setattr(subprocess, 'check_output', lambda _: SINFO_CONTENT.encode()) clusters = sinfo(username='john-doe', with_states=('completing'))