add tests for faiss package

drivendataorg · Jan 13, 2024 · dc65350 · dc65350
1 parent 9ba0c21
commit dc65350
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/runtime/tests/test_packages.py b/runtime/tests/test_packages.py
@@ -4,6 +4,7 @@
 import pytest
 
 packages = [
+    "faiss",
     "gensim",
     "keras",
     "numpy",
@@ -84,3 +85,10 @@ def test_spacy():
             ents.append(span)
         doc.ents = ents
         db.add(doc)
+
+
+@pytest.mark.skipif(not GPU_AVAILABLE, reason="No GPU available")
+def test_faiss_n_gpus():
+    import faiss
+    assert faiss.get_num_gpus() > 0
+