File size: 1,571 Bytes
fb12ddc | 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 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 | """Smoke test for multi-project vector store."""
import sys
import numpy as np
sys.path.insert(0, "HF_Space_hipVS")
from vector_store import get_store, list_projects, VectorStore
# Test multi-project isolation
DIM = 768 # CLIP dim for CPU
store_a = get_store("project-alpha", "image_index")
store_b = get_store("project-beta", "image_index")
vecs_a = np.random.randn(30, DIM).astype(np.float32)
vecs_b = np.random.randn(50, DIM).astype(np.float32)
store_a.add(vecs_a, [f"alpha_{i}" for i in range(30)])
store_b.add(vecs_b, [f"beta_{i}" for i in range(50)])
print(f"Store A: {store_a}")
print(f"Store B: {store_b}")
# Search in A should not return B's vectors
query = np.random.randn(DIM).astype(np.float32)
results_a = store_a.search(query, top_k=3)
results_b = store_b.search(query, top_k=3)
print(f"Search A: {[r['id'] for r in results_a]}")
print(f"Search B: {[r['id'] for r in results_b]}")
# Verify isolation
assert all("alpha" in r["id"] for r in results_a), "Project A returned non-alpha results!"
assert all("beta" in r["id"] for r in results_b), "Project B returned non-beta results!"
# Test append_and_rebuild
store_a.append_and_rebuild(np.random.randn(DIM).astype(np.float32), "alpha_new", {"test": True})
print(f"After append_and_rebuild: {store_a}")
# Test persistence
store_c = get_store("project-alpha", "image_index") # should be cached
assert store_c.count == 31
print(f"Cached store same ref: {store_c is store_a}")
# List projects
print(f"Projects: {list_projects()}")
# Cleanup
store_a.clear()
store_b.clear()
print("All tests passed")
|