fix test

Author: benoit-cty

tests/test_gpu.py

@@ -56,6 +56,18 @@ def setup_method(self):
         ):
             sys.modules.pop(module_name, None)
 
+        import codecarbon.core
+
+        for attr in ["gpu", "gpu_nvidia", "gpu_amd", "gpu_device"]:
+            if hasattr(codecarbon.core, attr):
+                delattr(codecarbon.core, attr)
+
+        import codecarbon.core
+
+        for attr in ["gpu", "gpu_nvidia", "gpu_amd", "gpu_device"]:
+            if hasattr(codecarbon.core, attr):
+                delattr(codecarbon.core, attr)
+
         # Setup the state, strings are returned as bytes
         self.DETAILS = {
             "handle_0": {
@@ -189,14 +201,18 @@ def test_gpu_no_power_limit(self):
         def raiseException(handle):
             raise Exception("Some bad exception")
 
-        pynvml.nvmlDeviceGetEnforcedPowerLimit = raiseException
-        alldevices = AllGPUDevices()
+        original_limit = pynvml.nvmlDeviceGetEnforcedPowerLimit
+        try:
+            pynvml.nvmlDeviceGetEnforcedPowerLimit = raiseException
+            alldevices = AllGPUDevices()
 
-        expected_power_limit = deepcopy(self.expected)
-        expected_power_limit[0]["power_limit"] = None
-        expected_power_limit[1]["power_limit"] = None
+            expected_power_limit = deepcopy(self.expected)
+            expected_power_limit[0]["power_limit"] = None
+            expected_power_limit[1]["power_limit"] = None
 
-        assert alldevices.get_gpu_details() == expected_power_limit
+            assert alldevices.get_gpu_details() == expected_power_limit
+        finally:
+            pynvml.nvmlDeviceGetEnforcedPowerLimit = original_limit
 
     def test_gpu_not_ready(self):
         import pynvml
@@ -206,14 +222,18 @@ def test_gpu_not_ready(self):
         def raise_exception(handle):
             raise pynvml.NVMLError("System is not in ready state")
 
-        pynvml.nvmlDeviceGetTotalEnergyConsumption = raise_exception
-        alldevices = AllGPUDevices()
+        original_energy = pynvml.nvmlDeviceGetTotalEnergyConsumption
+        try:
+            pynvml.nvmlDeviceGetTotalEnergyConsumption = raise_exception
+            alldevices = AllGPUDevices()
 
-        expected = deepcopy(self.expected)
-        expected[0]["total_energy_consumption"] = None
-        expected[1]["total_energy_consumption"] = None
+            expected = deepcopy(self.expected)
+            expected[0]["total_energy_consumption"] = None
+            expected[1]["total_energy_consumption"] = None
 
-        assert alldevices.get_gpu_details() == expected
+            assert alldevices.get_gpu_details() == expected
+        finally:
+            pynvml.nvmlDeviceGetTotalEnergyConsumption = original_energy
 
     def test_gpu_metadata_total_power(self):
         """
@@ -374,6 +394,18 @@ def setup_method(self):
         ):
             sys.modules.pop(module_name, None)
 
+        import codecarbon.core
+
+        for attr in ["gpu", "gpu_nvidia", "gpu_amd", "gpu_device"]:
+            if hasattr(codecarbon.core, attr):
+                delattr(codecarbon.core, attr)
+
+        import codecarbon.core
+
+        for attr in ["gpu", "gpu_nvidia", "gpu_amd", "gpu_device"]:
+            if hasattr(codecarbon.core, attr):
+                delattr(codecarbon.core, attr)
+
         import pynvml
 
         pynvml.INIT_MOCK.side_effect = pynvml.NVMLError("NVML Shared Library Not Found")
@@ -892,6 +924,16 @@ def _exec_gpu_module(self, import_func, check_output):
             "codecarbon.core.gpu_device",
         ):
             sys.modules.pop(module_name, None)
+        import codecarbon.core
+
+        for attr in ["gpu", "gpu_nvidia", "gpu_amd", "gpu_device"]:
+            if hasattr(codecarbon.core, attr):
+                delattr(codecarbon.core, attr)
+        import codecarbon.core
+
+        for attr in ["gpu", "gpu_nvidia", "gpu_amd", "gpu_device"]:
+            if hasattr(codecarbon.core, attr):
+                delattr(codecarbon.core, attr)
         with (
             mock.patch("subprocess.check_output", side_effect=check_output),
             mock.patch.object(builtins, "__import__", new=import_func),
@@ -1035,6 +1077,16 @@ def test_is_nvidia_system_fail(self, mock_subprocess):
 
 
 class TestGpuTracking:
+    def setup_method(self):
+        for module_name in list(sys.modules.keys()):
+            if module_name.startswith("codecarbon.core.gpu"):
+                sys.modules.pop(module_name, None)
+        import codecarbon.core
+
+        for attr in ["gpu", "gpu_nvidia", "gpu_amd", "gpu_device"]:
+            if hasattr(codecarbon.core, attr):
+                delattr(codecarbon.core, attr)
+
     @mock.patch("codecarbon.core.gpu.is_rocm_system", return_value=True)
     @mock.patch("codecarbon.core.gpu.is_nvidia_system", return_value=False)
     @mock.patch("codecarbon.core.gpu_amd.subprocess.check_output")