Skip to content

Commit

Permalink
add python tests
Browse files Browse the repository at this point in the history
  • Loading branch information
1yefuwang1 committed Jun 20, 2024
1 parent 7697a2b commit 41058ad
Show file tree
Hide file tree
Showing 4 changed files with 48 additions and 1 deletion.
Empty file.
43 changes: 43 additions & 0 deletions integration_test/python/test/vectorlite_test.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,43 @@
import vectorlite_py
import apsw
import pytest
import numpy as np


@pytest.fixture(scope='module')
def conn() -> None:
conn = apsw.Connection(':memory:')
conn.enable_load_extension(True)
conn.load_extension(vectorlite_py.vectorlite_path())
return conn

DIM = 32
NUM_ELEMENTS = 1000

# Generating sample data
@pytest.fixture(scope='module')
def random_vectors():
return np.float32(np.random.random((NUM_ELEMENTS, DIM)))

def test_vectorlite_info(conn):
cur = conn.cursor()
cur.execute('select vectorlite_info()')
output = cur.fetchone()
assert f'vectorlite extension version {vectorlite_py.__version__}' in output[0]

def test_l2_space_with_default_hnsw_param(conn, random_vectors):
cur = conn.cursor()
cur.execute('create virtual table x using vectorlite(my_embedding(32, "l2"), hnsw(max_elements=1000))')

for i in range(NUM_ELEMENTS):
cur.execute('insert into x (rowid, my_embedding) values (?, ?)', (i, random_vectors[i].tobytes()))

result = cur.execute('select my_embedding from x where rowid = 0').fetchone()
assert result[0] == random_vectors[0].tobytes()

cur.execute('delete from x where rowid = 0')
result = cur.execute('select my_embedding from x where rowid = 0').fetchone()
assert result is None

result = cur.execute('select rowid, distance from x where knn_search(my_embedding, knn_param(?, ?))', (random_vectors[0].tobytes(), 10)).fetchall()
assert len(result) == 10
4 changes: 3 additions & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,9 @@ requires = ["setuptools>=59", "wheel", "cmake", "ninja"]
build-backend = "setuptools.build_meta"

[tool.cibuildwheel]
test-requires = ["pytest", "numpy", "apsw>=3.42"]
test-command = "pytest {project}/integration_test/python/test"
skip = ["*-win32", "*-manylinux_i686", "*musllinux*", "pp*"]

[tool.cibuildwheel.macos]
skip = ["cp38*", "pp*"]
skip = ["cp36*", "cp37*", "cp38*", "pp*"]
2 changes: 2 additions & 0 deletions vectorlite_py/__init__.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
import os

__version__ = '0.1.0'

def vectorlite_path():
loadable_path = os.path.join(os.path.dirname(__file__), 'libvectorlite')
return os.path.normpath(loadable_path)
Expand Down

0 comments on commit 41058ad

Please sign in to comment.