fix: add tests for psutil availability and physical cpu count

Author: Arno

codecarbon/core/util.py

@@ -121,26 +121,37 @@ def count_physical_cpus():
             )
             return 1
 
+
 def _windows_get_physical_sockets():
     try:
         # use PowerShell to count number of objects of class Win32_Processor
-        cmd = ["powershell", "-NoProfile", "-Command", "(Get-CimInstance -ClassName Win32_Processor).Count"]      
-        result = subprocess.run(cmd, capture_output=True, text=True, timeout=10, check=True)
+        cmd = [
+            "powershell",
+            "-NoProfile",
+            "-Command",
+            "(Get-CimInstance -ClassName Win32_Processor).Count",
+        ]
+        result = subprocess.run(
+            cmd, capture_output=True, text=True, timeout=10, check=True
+        )
 
         output = result.stdout.strip()
 
-        # Fallback: if Count is empty, at least one socket exists
+        # Fallback: if empty, at least one socket exists
         if not output:
-            logger.debug("PowerShell command returned empty output for socket count. Defaulting to 1.")
-            output = 1           
+            logger.debug(
+                "PowerShell command returned empty output for socket count. Defaulting to 1."
+            )
+            output = 1
 
         logger.debug(f"Detected {output} physical sockets on Windows.")
         return int(output)
 
     except (subprocess.CalledProcessError, subprocess.TimeoutExpired, ValueError) as e:
         logger.error(f"Error detecting physical sockets on Windows: {e}")
         return 1  # Fallback:at least one socket
-    
+
+
 def count_cpus() -> int:
     if SLURM_JOB_ID is None:
         return psutil.cpu_count()

tests/test_cpu.py

@@ -11,13 +11,28 @@
     IntelPowerGadget,
     IntelRAPL,
     is_powergadget_available,
+    is_psutil_available,
 )
 from codecarbon.core.units import Energy, Power, Time
 from codecarbon.core.util import count_physical_cpus
 from codecarbon.external.hardware import CPU
 from codecarbon.input import DataSource
 
 
+class TestCPU(unittest.TestCase):
+    def test_is_psutil_available(self):
+        @mock.patch("psutil.cpu_times")
+        def test_is_psutil_available(self, mock_cpu_times):
+            # Test when psutil is available
+            mock_cpu_times.return_value = mock.Mock()
+            self.assertTrue(is_psutil_available())  # noqa: E712
+
+        @mock.patch("psutil.cpu_times", side_effect=AttributeError)
+        def test_is_psutil_not_available(self, mock_cpu_times):
+            # Test when psutil is not available
+            self.assertFalse(is_psutil_available())  # noqa: E712
+
+
 class TestIntelPowerGadget(unittest.TestCase):
     @pytest.mark.integ_test
     def test_intel_power_gadget(self):
@@ -306,10 +321,15 @@ def test_get_matching_cpu(self):
 class TestPhysicalCPU(unittest.TestCase):
     def test_count_physical_cpus_windows(self):
         with mock.patch("platform.system", return_value="Windows"):
-            with mock.patch.dict(os.environ, {"NUMBER_OF_PROCESSORS": "4"}):
+
+            with mock.patch(
+                "subprocess.run", return_value=mock.Mock(returncode=0, stdout="4")
+            ):
                 assert count_physical_cpus() == 4
 
-            with mock.patch.dict(os.environ, {}, clear=True):
+            with mock.patch(
+                "subprocess.run", return_value=mock.Mock(returncode=0, stdout="")
+            ):
                 assert count_physical_cpus() == 1
 
     def test_count_physical_cpus_linux(self):