update assert

ModelCloud · Dec 3, 2024 · b97f012 · b97f012
1 parent 112cbef
commit b97f012
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 7 deletions.
diff --git a/device_smi/base.py b/device_smi/base.py
@@ -36,7 +36,7 @@ def __init__(
     ):
         self.memory_used = memory_used
         self.memory_process = memory_process
-        self.utilization = utilization
+        self.utilization = max(0, utilization)
 
     def __str__(self):
         return str(self.__dict__)

diff --git a/tests/gpu.py b/tests/gpu.py
@@ -3,10 +3,13 @@
 dev = Device("gpu")
 print(dev)
 
-assert dev.type == "gpu"
-assert dev.pcie
-assert dev.gpu
+assert dev.type == "gpu", f"wrong type: {dev.type}"
+if dev.pcie:
+    assert dev.pcie.gen
+    assert dev.pcie.speed
+    assert dev.pcie.id
+if dev.gpu:
+    assert dev.gpu.driver
+    assert dev.gpu.firmware
 assert dev.model
-assert dev.memory_total > 10
-assert dev.vendor
-assert dev.features
+assert dev.memory_total > 10, f"wrong memory size: {dev.memory_total}"