Fix crc-idle casting issue for free memory when gpu resource is in dr…

…ain state (#261)

* Fixing the casting issue for drain gpu resources

* Fixing the test

* Adding drain case in GPU as a separate test

---------

Co-authored-by: Nickolas Comeau <[email protected]>

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:])

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"""