diff --git a/src/glicid_spawner/slurm.py b/src/glicid_spawner/slurm.py index 6a29a94..bdc1029 100644 --- a/src/glicid_spawner/slurm.py +++ b/src/glicid_spawner/slurm.py @@ -79,6 +79,9 @@ class SlurmPartition: def __iter__(self): return iter(self.nodes) + def __len__(self): + return len(self.nodes) + @property def gpus(self) -> str: """List of GPUs available.""" @@ -108,6 +111,9 @@ class SlurmCluster: def __iter__(self): return iter(self.partitions) + def __len__(self): + return len(self.partitions) + def __eq__(self, other): return str(self) == str(other) diff --git a/tests/test_slurm.py b/tests/test_slurm.py index 7901cc3..e93070f 100644 --- a/tests/test_slurm.py +++ b/tests/test_slurm.py @@ -73,6 +73,8 @@ def test_slurm_dataclasses(): assert str(partition) == 'standard' # = name assert partition.name == 'standard' + assert len(partition) == 1 + for _node in partition: assert str(_node) == 'cnode001' @@ -87,6 +89,8 @@ def test_slurm_dataclasses(): assert cluster.name == 'nautilus' assert cluster == 'nautilus' # __eq__ + assert len(cluster) == 1 + for _partition in cluster: assert str(_partition) == 'standard'