From 8415b3308d1673e4122c9aed9dda7e96d96e3826 Mon Sep 17 00:00:00 2001 From: jeffnvidia Date: Mon, 27 May 2024 17:53:48 +0300 Subject: [PATCH] add tests --- tests/test_slurm_system.py | 97 +++++++++++++++++++++++++++++++------- 1 file changed, 81 insertions(+), 16 deletions(-) diff --git a/tests/test_slurm_system.py b/tests/test_slurm_system.py index 5b98e9b96..e19467c1b 100644 --- a/tests/test_slurm_system.py +++ b/tests/test_slurm_system.py @@ -8,15 +8,31 @@ @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), + ] + backup_nodes = [ + SlurmNode(name="node01", partition="backup", state=SlurmNodeState.UNKNOWN_STATE), + SlurmNode(name="node02", partition="backup", state=SlurmNodeState.UNKNOWN_STATE), + SlurmNode(name="node03", partition="backup", state=SlurmNodeState.UNKNOWN_STATE), + SlurmNode(name="node04", partition="backup", state=SlurmNodeState.UNKNOWN_STATE), + SlurmNode(name="node05", partition="backup", state=SlurmNodeState.UNKNOWN_STATE), + SlurmNode(name="node06", partition="backup", state=SlurmNodeState.UNKNOWN_STATE), + SlurmNode(name="node07", partition="backup", state=SlurmNodeState.UNKNOWN_STATE), + SlurmNode(name="node08", partition="backup", state=SlurmNodeState.UNKNOWN_STATE), ] system = SlurmSystem( name="test_system", install_path="/fake/path", output_path="/fake/output", default_partition="main", - partitions={"main": nodes}, + partitions={"main": nodes, "backup": backup_nodes}, ) return system @@ -29,7 +45,7 @@ def test_parse_squeue_output(slurm_system): def test_parse_squeue_output_with_node_ranges_and_root_user(slurm_system): - squeue_output = "nodeA[001-008],nodeB[001-008]|root" + squeue_output = "nodeA[001-008]|root\nnodeB[001-008]|root" user_map = slurm_system.parse_squeue_output(squeue_output) expected_nodes = [f"nodeA{str(i).zfill(3)}" for i in range(1, 9)] + [f"nodeB{str(i).zfill(3)}" for i in range(1, 9)] @@ -39,24 +55,73 @@ def test_parse_squeue_output_with_node_ranges_and_root_user(slurm_system): 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 idle node[01-16] + """ + node_user_map = { + "": "user1", + "node-033": "user2", + "node-034": "user3", + "node-056": "user3", + "node-057": "user3", + "node-058": "user3", + "node-049": "user4", + "node-050": "user4", + "node01": "user5", + "node02": "user5", + "node03": "user5", + "node04": "user5", + "node05": "user5", + "node06": "user5", + "node07": "user5", + "node08": "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 + for i in range(len(slurm_system.partitions["main"])): + assert slurm_system.partitions["main"][i].state == SlurmNodeState.ALLOCATED + for i in range(len(slurm_system.partitions["backup"])): + assert slurm_system.partitions["backup"][i].state == SlurmNodeState.IDLE @patch("cloudai.schema.system.SlurmSystem.get_squeue") @patch("cloudai.schema.system.SlurmSystem.get_sinfo") def test_update_node_states_with_mocked_outputs(mock_get_sinfo, mock_get_squeue, slurm_system): - mock_get_squeue.return_value = "nodeA001|root" - mock_get_sinfo.return_value = "PARTITION AVAIL TIMELIMIT NODES STATE NODELIST\n" "main up infinite 1 idle nodeA001" + mock_get_squeue.return_value = "node-115|user1" + mock_get_sinfo.return_value = "PARTITION AVAIL TIMELIMIT NODES STATE NODELIST\n" "main up infinite 1 idle node-115" + + slurm_system.update_node_states() + for node in slurm_system.partitions["main"]: + if node.name == "node-115": + assert node.state == SlurmNodeState.IDLE + assert node.user == "user1" + + mock_get_squeue.return_value = "node01|root" + mock_get_sinfo.return_value = ( + "PARTITION AVAIL TIMELIMIT NODES STATE NODELIST\n" "backup up infinite 1 allocated node01" + ) slurm_system.update_node_states() + for node in slurm_system.partitions["backup"]: + if node.name == "node01": + assert node.state == SlurmNodeState.ALLOCATED + assert node.user == "root" + + +def test_parse_node_list(slurm_system): + node_lists = ["node-[048-051]", "node-[055,114]", "", "node-001", "node[1-4]"] + + expected_outputs = [ + ["node-048", "node-049", "node-050", "node-051"], + ["node-055", "node-114"], + [], + ["node-001"], + ["node1", "node2", "node3", "node4"], + ] - assert slurm_system.partitions["main"][0].state == SlurmNodeState.IDLE - assert slurm_system.partitions["main"][0].user == "root" + for i, node_list in enumerate(node_lists): + parsed_nodes = slurm_system.parse_node_list(node_list) + assert parsed_nodes == expected_outputs[i]