| 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") |
|
|
