File size: 1,256 Bytes
1d1ee9b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 | import numpy as np
import pytest
from common.vector_index import create_index, save_index, load_index, add_vector, search
def test_add_vector_returns_sequential_ids():
index = create_index(dim=4)
v1 = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32)
v2 = np.array([0.0, 1.0, 0.0, 0.0], dtype=np.float32)
assert add_vector(index, v1) == 0
assert add_vector(index, v2) == 1
assert index.ntotal == 2
def test_search_returns_closest_vector_first():
index = create_index(dim=4)
add_vector(index, np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32))
add_vector(index, np.array([0.0, 1.0, 0.0, 0.0], dtype=np.float32))
query = np.array([0.9, 0.1, 0.0, 0.0], dtype=np.float32)
results = search(index, query, k=2)
assert results[0][0] == 0
assert results[0][1] > results[1][1]
def test_save_and_load_roundtrip(tmp_path):
index = create_index(dim=4)
add_vector(index, np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32))
path = str(tmp_path / "index.faiss")
save_index(index, path)
loaded = load_index(path)
assert loaded.ntotal == 1
assert loaded.d == 4
def test_load_missing_index_raises():
with pytest.raises(FileNotFoundError):
load_index("/nonexistent/path/index.faiss")
|