Skip to content

Commit 53a5db3

Browse files
committed
add python tests
1 parent 7697a2b commit 53a5db3

File tree

4 files changed

+48
-1
lines changed

4 files changed

+48
-1
lines changed

integration_test/python/test/__init__.py

Whitespace-only changes.
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
import vectorlite_py
2+
import apsw
3+
import pytest
4+
import numpy as np
5+
6+
7+
@pytest.fixture(scope='module')
8+
def conn() -> None:
9+
conn = apsw.Connection(':memory:')
10+
conn.enable_load_extension(True)
11+
conn.load_extension(vectorlite_py.vectorlite_path())
12+
return conn
13+
14+
DIM = 32
15+
NUM_ELEMENTS = 1000
16+
17+
# Generating sample data
18+
@pytest.fixture(scope='module')
19+
def random_vectors():
20+
return np.float32(np.random.random((NUM_ELEMENTS, DIM)))
21+
22+
def test_vectorlite_info(conn):
23+
cur = conn.cursor()
24+
cur.execute('select vectorlite_info()')
25+
output = cur.fetchone()
26+
assert f'vectorlite extension version {vectorlite_py.__version__}' in output[0]
27+
28+
def test_l2_space_with_default_hnsw_param(conn, random_vectors):
29+
cur = conn.cursor()
30+
cur.execute('create virtual table x using vectorlite(my_embedding(32, "l2"), hnsw(max_elements=1000))')
31+
32+
for i in range(NUM_ELEMENTS):
33+
cur.execute('insert into x (rowid, my_embedding) values (?, ?)', (i, random_vectors[i].tobytes()))
34+
35+
result = cur.execute('select my_embedding from x where rowid = 0').fetchone()
36+
assert result[0] == random_vectors[0].tobytes()
37+
38+
cur.execute('delete from x where rowid = 0')
39+
result = cur.execute('select my_embedding from x where rowid = 0').fetchone()
40+
assert result is None
41+
42+
result = cur.execute('select rowid, distance from x where knn_search(my_embedding, knn_param(?, ?))', (random_vectors[0].tobytes(), 10)).fetchall()
43+
assert len(result) == 10

pyproject.toml

+3-1
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,9 @@ requires = ["setuptools>=59", "wheel", "cmake", "ninja"]
44
build-backend = "setuptools.build_meta"
55

66
[tool.cibuildwheel]
7+
test-requires = ["pytest", "numpy", "apsw>=3.46"]
8+
test-command = "pytest {project}/integration_test/python/test"
79
skip = ["*-win32", "*-manylinux_i686", "*musllinux*", "pp*"]
810

911
[tool.cibuildwheel.macos]
10-
skip = ["cp38*", "pp*"]
12+
skip = ["cp36*", "cp37*", "cp38*", "pp*"]

vectorlite_py/__init__.py

+2
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
import os
22

3+
__version__ = '0.1.0'
4+
35
def vectorlite_path():
46
loadable_path = os.path.join(os.path.dirname(__file__), 'libvectorlite')
57
return os.path.normpath(loadable_path)

0 commit comments

Comments
 (0)