From 02166be4d23b6c94d6e9195c65681f0f30ae14d1 Mon Sep 17 00:00:00 2001
From: Benoit Seignovert <benoit.seignovert@univ-nantes.fr>
Date: Mon, 19 Feb 2024 17:53:35 +0100
Subject: [PATCH] Add getitem and eq to slurm objects

---
 src/glicid_spawner/slurm.py | 20 ++++++++++++++++----
 tests/test_slurm.py         | 31 +++++++++++++++----------------
 2 files changed, 31 insertions(+), 20 deletions(-)

diff --git a/src/glicid_spawner/slurm.py b/src/glicid_spawner/slurm.py
index 8c8300a..c059b0e 100644
--- a/src/glicid_spawner/slurm.py
+++ b/src/glicid_spawner/slurm.py
@@ -40,6 +40,9 @@ class SlurmGpu:
     def __str__(self):
         return self.name
 
+    def __eq__(self, other):
+        return str(self) == str(other)
+
 
 @dataclass
 class SlurmNode:
@@ -65,6 +68,9 @@ class SlurmNode:
     def __str__(self):
         return self.hostname
 
+    def __eq__(self, other):
+        return str(self) == str(other)
+
 
 @dataclass
 class SlurmPartition:
@@ -76,14 +82,17 @@ class SlurmPartition:
     def __str__(self):
         return self.name
 
+    def __eq__(self, other):
+        return str(self) == str(other)
+
     def __iter__(self):
         return iter(self.nodes)
 
     def __len__(self):
         return len(self.nodes)
 
-    def __eq__(self, other):
-        return str(self) == str(other)
+    def __getitem__(self, other):
+        return self.nodes[self.nodes.index(other)]
 
     @property
     def gpus(self) -> str:
@@ -111,14 +120,17 @@ class SlurmCluster:
     def __str__(self):
         return self.name
 
+    def __eq__(self, other):
+        return str(self) == str(other)
+
     def __iter__(self):
         return iter(self.partitions)
 
     def __len__(self):
         return len(self.partitions)
 
-    def __eq__(self, other):
-        return str(self) == str(other)
+    def __getitem__(self, other):
+        return self.partitions[self.partitions.index(other)]
 
 
 def sinfo_run(username: str = None) -> str:
diff --git a/tests/test_slurm.py b/tests/test_slurm.py
index e5aba42..3f8185c 100644
--- a/tests/test_slurm.py
+++ b/tests/test_slurm.py
@@ -38,8 +38,7 @@ def test_slurm_dataclasses():
     gpu = SlurmGpu('fOo', '1')
 
     assert gpu  # __bool__
-    assert str(gpu) == 'Foo'  # = name
-    assert gpu.name == 'Foo'
+    assert gpu == str(gpu) == gpu.name == 'Foo'
     assert gpu.nb == 1
 
     assert isinstance(gpu.name, str)
@@ -49,17 +48,15 @@ def test_slurm_dataclasses():
     gpu = SlurmGpu()
 
     assert not gpu  # __bool__
-    assert str(gpu) == 'None'  # = name
-    assert gpu.name == 'None'
+    assert str(gpu) == gpu.name == 'None'
     assert gpu.nb == 0
 
     # Node
     node = SlurmNode(*'nautilus standard cnode001 completing 0/96/0/96 384000 (null)'.split())
 
-    assert str(node) == 'cnode001'  # hostname
     assert node.cluster == 'nautilus'
     assert node.partition == 'standard'
-    assert node.hostname == 'cnode001'
+    assert node == str(node) == node.hostname == 'cnode001'
     assert node.state == 'completing'
     assert node.cpu.allocated == 0
     assert node.cpu.idle == 96
@@ -70,14 +67,15 @@ def test_slurm_dataclasses():
     # Partition
     partition = SlurmPartition('standard', [node])
 
-    assert str(partition) == 'standard'  # = name
-    assert partition.name == 'standard'
-    assert partition == 'standard'  # __eq__ on name
-
+    assert partition == str(partition) == partition.name == 'standard'
     assert len(partition) == 1
 
     for _node in partition:
-        assert str(_node) == 'cnode001'
+        assert _node == 'cnode001'
+
+    _node = partition['cnode001']  # __getitem__
+    assert isinstance(_node, SlurmNode)
+    assert _node == 'cnode001'
 
     assert partition.gpus == 'None'
     assert partition.max_idle_cpu == 96
@@ -86,14 +84,15 @@ def test_slurm_dataclasses():
     # Cluster
     cluster = SlurmCluster('nautilus', [partition])
 
-    assert str(cluster) == 'nautilus'  # = name
-    assert cluster.name == 'nautilus'
-    assert cluster == 'nautilus'  # __eq__ on name
-
+    assert cluster == str(cluster) == cluster.name == 'nautilus'
     assert len(cluster) == 1
 
     for _partition in cluster:
-        assert str(_partition) == 'standard'
+        assert _partition == 'standard'
+
+    _partition = cluster['standard']  # __getitem__
+    assert isinstance(_partition, SlurmPartition)
+    assert _partition == 'standard'
 
 
 def test_slurm_sinfo_run(monkeypatch):