tests: Fix test_get_gpu_info

devitocodes · Feb 7, 2022 · 7bb6ff5 · 7bb6ff5
1 parent f49a856
commit 7bb6ff5
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/test_gpu_common.py b/tests/test_gpu_common.py
@@ -20,8 +20,9 @@ class TestGPUInfo(object):
 
     def test_get_gpu_info(self):
         info = get_gpu_info()
+        known = ['nvidia', 'tesla', 'geforce', 'unspecified']
         try:
-            assert info['architecture'].lower() in ['tesla', 'geforce', 'unspecified']
+            assert info['architecture'].lower() in known
         except KeyError:
             # There might be than one GPUs, but for now we don't care
             # as we're not really exploiting this info yet...