conftest can test without torch

lancedb · Dec 24, 2022 · 9e24247 · 9e24247
1 parent 854cd62
commit 9e24247
Showing 1 changed file with 10 additions and 2 deletions.
diff --git a/integration/duckdb/tests/conftest.py b/integration/duckdb/tests/conftest.py
@@ -4,15 +4,23 @@
 
 import duckdb
 import pytest
-import torch
+
+try:
+    import torch
+    has_torch = True
+except ImportError:
+    has_torch = False
 
 
 @pytest.fixture
 def db() -> duckdb.DuckDBPyConnection:
     """Initialize duckdb with lance extension"""
     db = duckdb.connect(config={"allow_unsigned_extensions": True})
 
-    build_dir = "cuda-build" if torch.cuda.is_available() else "manylinux-build"
+    build_dir = "manylinux-build"
+    if has_torch:
+        if torch.cuda.is_available():
+            build_dir = "cuda-build"
 
     cur_path = Path(__file__).parent
     db.install_extension(str(cur_path.parent / build_dir / "lance.duckdb_extension"), force_install=True)