diff --git a/src/glicid_spawner/slurm.py b/src/glicid_spawner/slurm.py index bdc1029..8c8300a 100644 --- a/src/glicid_spawner/slurm.py +++ b/src/glicid_spawner/slurm.py @@ -82,6 +82,9 @@ class SlurmPartition: def __len__(self): return len(self.nodes) + def __eq__(self, other): + return str(self) == str(other) + @property def gpus(self) -> str: """List of GPUs available.""" diff --git a/tests/test_slurm.py b/tests/test_slurm.py index e93070f..e5aba42 100644 --- a/tests/test_slurm.py +++ b/tests/test_slurm.py @@ -72,6 +72,7 @@ def test_slurm_dataclasses(): assert str(partition) == 'standard' # = name assert partition.name == 'standard' + assert partition == 'standard' # __eq__ on name assert len(partition) == 1 @@ -87,7 +88,7 @@ def test_slurm_dataclasses(): assert str(cluster) == 'nautilus' # = name assert cluster.name == 'nautilus' - assert cluster == 'nautilus' # __eq__ + assert cluster == 'nautilus' # __eq__ on name assert len(cluster) == 1