diff --git a/apps/crc_idle.py b/apps/crc_idle.py index 0d567db..e7c87f8 100755 --- a/apps/crc_idle.py +++ b/apps/crc_idle.py @@ -121,6 +121,7 @@ def _count_idle_gpu_resources(cluster: str, partition: str) -> dict[int, dict[st # If the node is in a downed state, report 0 resource availability. if re.search("drain", state): idle = 0 + free_mem = 0 else: allocated = int(allocated[-1:]) diff --git a/tests/test_crc_idle.py b/tests/test_crc_idle.py index 7122d87..3c89bd3 100644 --- a/tests/test_crc_idle.py +++ b/tests/test_crc_idle.py @@ -95,7 +95,7 @@ def test_count_idle_gpu_resources(self, mock_run_command: Mock) -> None: cluster = 'gpu' partition = 'default' - mock_run_command.return_value = "node1_4_2_idle_3500\nnode2_4_4_drain_4000" + mock_run_command.return_value = "node1_4_2_idle_3500\nnode2_4_4_alloc_4000" app = CrcIdle() result = app.count_idle_resources(cluster, partition) @@ -104,6 +104,18 @@ def test_count_idle_gpu_resources(self, mock_run_command: Mock) -> None: } self.assertEqual(expected, result) + @patch('apps.utils.Shell.run_command') + def test_count_drain_gpu_resources(self, mock_run_command: Mock) -> None: + """Test counting drain GPU resources.""" + + cluster = 'gpu' + partition = 'default' + mock_run_command.return_value = "node1_4_2_drain*_N/A\nnode2_4_4_drain_N/A" + + app = CrcIdle() + result = app.count_idle_resources(cluster, partition) + expected = {0: {'count': 2, 'min_free_mem': 0, 'max_free_mem': 0}} + self.assertEqual(expected, result) class PrintPartitionSummary(TestCase): """Test the printing of a partition summary"""