update assert

ModelCloud · Dec 3, 2024 · 06bc767 · 06bc767
1 parent 112cbef
commit 06bc767
Showing 1 changed file with 9 additions and 6 deletions.
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}"