From 41058adb021345f73b9b0c2596c0a6c141a356ce Mon Sep 17 00:00:00 2001
From: Yefu Wang <1yefuwang1@gmail.com>
Date: Thu, 20 Jun 2024 00:04:20 +0800
Subject: [PATCH] add python tests

---
 integration_test/python/test/__init__.py      |  0
 .../python/test/vectorlite_test.py            | 43 +++++++++++++++++++
 pyproject.toml                                |  4 +-
 vectorlite_py/__init__.py                     |  2 +
 4 files changed, 48 insertions(+), 1 deletion(-)
 create mode 100644 integration_test/python/test/__init__.py
 create mode 100644 integration_test/python/test/vectorlite_test.py

diff --git a/integration_test/python/test/__init__.py b/integration_test/python/test/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/integration_test/python/test/vectorlite_test.py b/integration_test/python/test/vectorlite_test.py
new file mode 100644
index 0000000..d564b64
--- /dev/null
+++ b/integration_test/python/test/vectorlite_test.py
@@ -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
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
index f231387..316689d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -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*"]
diff --git a/vectorlite_py/__init__.py b/vectorlite_py/__init__.py
index 80ffc1d..0352d3a 100644
--- a/vectorlite_py/__init__.py
+++ b/vectorlite_py/__init__.py
@@ -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)