add test for parse_sinfo

NVIDIA · May 23, 2024 · df011de · df011de
1 parent 0d35a8b
commit df011de
Showing 1 changed file with 20 additions and 12 deletions.
diff --git a/tests/test_slurm_system.py b/tests/test_slurm_system.py
@@ -8,8 +8,14 @@
 @pytest.fixture
 def slurm_system():
     nodes = [
-        SlurmNode(name="nodeA001", partition="main", state=SlurmNodeState.UNKNOWN_STATE),
-        SlurmNode(name="nodeB001", partition="main", state=SlurmNodeState.UNKNOWN_STATE),
+        SlurmNode(name="node-115", partition="main", state=SlurmNodeState.UNKNOWN_STATE),
+        SlurmNode(name="node-116", partition="main", state=SlurmNodeState.UNKNOWN_STATE),
+        SlurmNode(name="node-117", partition="main", state=SlurmNodeState.UNKNOWN_STATE),
+        SlurmNode(name="node-118", partition="main", state=SlurmNodeState.UNKNOWN_STATE),
+        SlurmNode(name="node-119", partition="main", state=SlurmNodeState.UNKNOWN_STATE),
+        SlurmNode(name="node-120", partition="main", state=SlurmNodeState.UNKNOWN_STATE),
+        SlurmNode(name="node-121", partition="main", state=SlurmNodeState.UNKNOWN_STATE),
+        SlurmNode(name="node-122", partition="main", state=SlurmNodeState.UNKNOWN_STATE),
     ]
     system = SlurmSystem(
         name="test_system",
@@ -36,18 +42,20 @@ def test_parse_squeue_output_with_node_ranges_and_root_user(slurm_system):
     expected_map = {node: "root" for node in expected_nodes}
 
     assert user_map == expected_map, "All nodes should be mapped to 'root'"
-
-
+
 def test_parse_sinfo_output(slurm_system):
-    sinfo_output = (
-        "PARTITION AVAIL TIMELIMIT NODES STATE NODELIST\n"
-        "main up infinite 1 idle nodeA001\n"
-        "main up infinite 1 idle nodeB001"
-    )
-    node_user_map = {"nodeA001": "root", "nodeB001": "user"}
+    sinfo_output = """
+    PARTITION AVAIL  TIMELIMIT  NODES  STATE NODELIST
+    main    up    3:00:00      1  inval node-081
+    main    up    3:00:00      5  drain node-[065-066,114,124-125]
+    main    up    3:00:00      2   resv node-[034-035]
+    main    up    3:00:00     88  alloc node-[033,036-064,067-080,082-113,115-123,126-128]
+    backup    up   12:00:00     16  alloc node-[01-16]
+    """
+    node_user_map = {'': 'user1', 'node-033': 'user2', 'node-[036-064': 'user3', '067-080': 'user3', '082-113': 'user3', '115-118]': 'user3', 'node-[119-123': 'user4', '126-128]': 'user4', 'node-01': 'user5', 'node-02': 'user5', 'node-03': 'user5', 'node-04': 'user5', 'node-05': 'user5', 'node-06': 'user5', 'node-07': 'user5', 'node-08': 'user5', 'node-09': 'user5', 'node-10': 'user5', 'node-11': 'user5', 'node-12': 'user5', 'node-13': 'user5', 'node-14': 'user5', 'node-15': 'user5', 'node-16': 'user5'}
     slurm_system.parse_sinfo_output(sinfo_output, node_user_map)
-    assert slurm_system.partitions["main"][0].state == SlurmNodeState.IDLE
-    assert slurm_system.partitions["main"][1].state == SlurmNodeState.IDLE
+    assert slurm_system.partitions["main"][0].state == SlurmNodeState.ALLOCATED
+    assert slurm_system.partitions["main"][1].state == SlurmNodeState.ALLOCATED
 
 
 @patch("cloudai.schema.system.SlurmSystem.get_squeue")