459 lines
17 KiB
Python
459 lines
17 KiB
Python
"""
|
|
FAISS向量管理器测试文件
|
|
用于验证FAISSManager的各项功能
|
|
"""
|
|
|
|
import unittest
|
|
import numpy as np
|
|
import tempfile
|
|
import shutil
|
|
import os
|
|
import sys
|
|
|
|
# 添加项目根目录到路径
|
|
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from faiss_vector_db.faiss_manager import FAISSManager
|
|
|
|
class TestFAISSManager(unittest.TestCase):
|
|
"""FAISS管理器测试类"""
|
|
|
|
def setUp(self):
|
|
"""测试前的设置"""
|
|
self.dimension = 64
|
|
self.manager = FAISSManager(
|
|
dimension=self.dimension,
|
|
index_type="IndexFlatIP",
|
|
normalize_vectors=True,
|
|
metric_type="cosine"
|
|
)
|
|
|
|
# 创建临时目录用于测试保存/加载
|
|
self.temp_dir = tempfile.mkdtemp()
|
|
|
|
def tearDown(self):
|
|
"""测试后的清理"""
|
|
# 删除临时目录
|
|
if os.path.exists(self.temp_dir):
|
|
shutil.rmtree(self.temp_dir)
|
|
|
|
def test_initialization(self):
|
|
"""测试初始化"""
|
|
self.assertEqual(self.manager.dimension, self.dimension)
|
|
self.assertEqual(self.manager.index_type, "IndexFlatIP")
|
|
self.assertTrue(self.manager.normalize_vectors)
|
|
self.assertEqual(self.manager.metric_type, "cosine")
|
|
self.assertEqual(self.manager.index.ntotal, 0)
|
|
|
|
def test_add_single_vector(self):
|
|
"""测试添加单个向量,用numpy随机生成一个向量"""
|
|
vector = np.random.random(self.dimension).astype(np.float32)
|
|
vector_id = "test_001"
|
|
metadata = {"name": "测试菜品", "category": "测试"}
|
|
|
|
# 测试添加
|
|
result = self.manager.add_vector(vector, vector_id, metadata)
|
|
self.assertTrue(result)
|
|
|
|
# 验证索引状态
|
|
self.assertEqual(self.manager.index.ntotal, 1)
|
|
self.assertIn(vector_id, self.manager.reverse_id_mapping)
|
|
self.assertIn(vector_id, self.manager.metadata)
|
|
|
|
# 测试重复添加
|
|
result = self.manager.add_vector(vector, vector_id, metadata)
|
|
self.assertTrue(result) # 应该更新成功
|
|
self.assertEqual(self.manager.index.ntotal, 2) # 因为，没有真正的删除
|
|
|
|
def test_add_batch_vectors(self):
|
|
"""测试批量添加向量"""
|
|
batch_size = 10
|
|
vectors = np.random.random((batch_size, self.dimension)).astype(np.float32)
|
|
vector_ids = [f"batch_{i}" for i in range(batch_size)]
|
|
metadata_list = [{"name": f"菜品{i}", "category": "测试"} for i in range(batch_size)]
|
|
|
|
# 测试批量添加
|
|
results = self.manager.add_vectors_batch(vectors, vector_ids, metadata_list)
|
|
|
|
# 验证结果
|
|
self.assertEqual(len(results), batch_size)
|
|
self.assertTrue(all(results))
|
|
self.assertEqual(self.manager.index.ntotal, batch_size)
|
|
|
|
# 验证所有向量都被添加
|
|
for vector_id in vector_ids:
|
|
self.assertIn(vector_id, self.manager.reverse_id_mapping)
|
|
self.assertIn(vector_id, self.manager.metadata)
|
|
|
|
def test_search_similar(self):
|
|
"""测试相似度搜索"""
|
|
# 先添加一些向量
|
|
batch_size = 20
|
|
vectors = np.random.random((batch_size, self.dimension)).astype(np.float32)
|
|
vector_ids = [f"search_test_{i}" for i in range(batch_size)]
|
|
metadata_list = [{"name": f"菜品{i}", "score": i} for i in range(batch_size)]
|
|
|
|
self.manager.add_vectors_batch(vectors, vector_ids, metadata_list)
|
|
|
|
# 测试搜索
|
|
query_vector = np.random.random(self.dimension).astype(np.float32)
|
|
k = 5
|
|
results = self.manager.search_similar(query_vector, k=k)
|
|
|
|
# 验证结果
|
|
self.assertLessEqual(len(results), k)
|
|
self.assertLessEqual(len(results), batch_size)
|
|
|
|
# 验证结果格式
|
|
for vector_id, similarity, metadata in results:
|
|
self.assertIsInstance(vector_id, str)
|
|
self.assertIsInstance(similarity, float)
|
|
self.assertIsInstance(metadata, dict)
|
|
self.assertIn(vector_id, self.manager.reverse_id_mapping)
|
|
|
|
# 测试带阈值的搜索
|
|
threshold = 0.8
|
|
results_with_threshold = self.manager.search_similar(
|
|
query_vector, k=k, threshold=threshold
|
|
)
|
|
|
|
# 验证阈值过滤
|
|
for _, similarity, _ in results_with_threshold:
|
|
self.assertGreaterEqual(similarity, threshold)
|
|
|
|
def test_get_vector_by_id(self):
|
|
"""测试根据ID获取向量"""
|
|
vector = np.random.random(self.dimension).astype(np.float32)
|
|
vector_id = "get_test_001"
|
|
metadata = {"name": "测试获取", "value": 42}
|
|
|
|
# 添加向量
|
|
self.manager.add_vector(vector, vector_id, metadata)
|
|
|
|
# 测试获取存在的向量
|
|
result = self.manager.get_vector_by_id(vector_id)
|
|
self.assertIsNotNone(result)
|
|
vector_data, retrieved_metadata = result
|
|
self.assertEqual(retrieved_metadata, metadata)
|
|
|
|
# 测试获取不存在的向量
|
|
result = self.manager.get_vector_by_id("nonexistent")
|
|
self.assertIsNone(result)
|
|
|
|
def test_update_vector(self):
|
|
"""测试更新向量"""
|
|
# 添加初始向量
|
|
vector = np.random.random(self.dimension).astype(np.float32)
|
|
vector_id = "update_test_001"
|
|
metadata = {"name": "原始菜品", "version": 1}
|
|
|
|
self.manager.add_vector(vector, vector_id, metadata)
|
|
initial_count = self.manager.index.ntotal
|
|
|
|
# 更新向量
|
|
new_vector = np.random.random(self.dimension).astype(np.float32)
|
|
new_metadata = {"name": "更新菜品", "version": 2}
|
|
|
|
result = self.manager.update_vector(new_vector, vector_id, new_metadata)
|
|
self.assertTrue(result)
|
|
|
|
# 验证更新后的状态
|
|
self.assertEqual(self.manager.index.ntotal, initial_count)
|
|
|
|
# 验证元数据已更新
|
|
vector_info = self.manager.get_vector_by_id(vector_id)
|
|
self.assertIsNotNone(vector_info)
|
|
_, retrieved_metadata = vector_info
|
|
self.assertEqual(retrieved_metadata, new_metadata)
|
|
|
|
# 测试更新不存在的向量
|
|
result = self.manager.update_vector(new_vector, "nonexistent", new_metadata)
|
|
self.assertFalse(result)
|
|
|
|
def test_delete_vector(self):
|
|
"""测试删除向量"""
|
|
# 添加向量
|
|
vector = np.random.random(self.dimension).astype(np.float32)
|
|
vector_id = "delete_test_001"
|
|
metadata = {"name": "待删除菜品"}
|
|
|
|
self.manager.add_vector(vector, vector_id, metadata)
|
|
|
|
# 验证向量存在
|
|
self.assertIn(vector_id, self.manager.reverse_id_mapping)
|
|
self.assertIn(vector_id, self.manager.metadata)
|
|
|
|
# 删除向量
|
|
result = self.manager.delete_vector(vector_id)
|
|
self.assertTrue(result)
|
|
|
|
# 验证向量已删除
|
|
self.assertNotIn(vector_id, self.manager.reverse_id_mapping)
|
|
self.assertNotIn(vector_id, self.manager.metadata)
|
|
|
|
# 测试删除不存在的向量
|
|
result = self.manager.delete_vector("nonexistent")
|
|
self.assertFalse(result)
|
|
|
|
def test_delete_vectors_batch(self):
|
|
"""测试批量删除向量"""
|
|
# 添加多个向量
|
|
batch_size = 5
|
|
vectors = np.random.random((batch_size, self.dimension)).astype(np.float32)
|
|
vector_ids = [f"batch_delete_{i}" for i in range(batch_size)]
|
|
|
|
self.manager.add_vectors_batch(vectors, vector_ids)
|
|
|
|
# 批量删除部分向量
|
|
delete_ids = vector_ids[:3]
|
|
results = self.manager.delete_vectors_batch(delete_ids)
|
|
|
|
# 验证删除结果
|
|
self.assertEqual(len(results), len(delete_ids))
|
|
self.assertTrue(all(results))
|
|
|
|
# 验证向量已删除
|
|
for vector_id in delete_ids:
|
|
self.assertNotIn(vector_id, self.manager.reverse_id_mapping)
|
|
|
|
# 验证剩余向量仍存在
|
|
for vector_id in vector_ids[3:]:
|
|
self.assertIn(vector_id, self.manager.reverse_id_mapping)
|
|
|
|
def test_save_and_load_index(self):
|
|
"""测试保存和加载索引"""
|
|
# 添加一些向量
|
|
batch_size = 10
|
|
vectors = np.random.random((batch_size, self.dimension)).astype(np.float32)
|
|
vector_ids = [f"save_load_{i}" for i in range(batch_size)]
|
|
metadata_list = [{"name": f"菜品{i}", "id": i} for i in range(batch_size)]
|
|
|
|
self.manager.add_vectors_batch(vectors, vector_ids, metadata_list)
|
|
|
|
# 保存索引
|
|
save_path = os.path.join(self.temp_dir, "test_index")
|
|
result = self.manager.save_index(save_path)
|
|
self.assertTrue(result)
|
|
|
|
# 验证文件已创建
|
|
self.assertTrue(os.path.exists(os.path.join(save_path, "faiss_index.bin")))
|
|
self.assertTrue(os.path.exists(os.path.join(save_path, "metadata.pkl")))
|
|
self.assertTrue(os.path.exists(os.path.join(save_path, "config.json")))
|
|
|
|
# 创建新管理器并加载索引
|
|
new_manager = FAISSManager(dimension=self.dimension)
|
|
result = new_manager.load_index(save_path)
|
|
self.assertTrue(result)
|
|
|
|
# 验证加载后的状态
|
|
self.assertEqual(new_manager.index.ntotal, batch_size)
|
|
self.assertEqual(len(new_manager.reverse_id_mapping), batch_size)
|
|
self.assertEqual(len(new_manager.metadata), batch_size)
|
|
|
|
# 验证数据一致性
|
|
for vector_id in vector_ids:
|
|
self.assertIn(vector_id, new_manager.reverse_id_mapping)
|
|
self.assertIn(vector_id, new_manager.metadata)
|
|
|
|
# 测试加载后的搜索功能
|
|
query_vector = np.random.random(self.dimension).astype(np.float32)
|
|
results = new_manager.search_similar(query_vector, k=3)
|
|
self.assertGreater(len(results), 0)
|
|
|
|
def test_clear_index(self):
|
|
"""测试清空索引"""
|
|
# 添加一些向量
|
|
vectors = np.random.random((5, self.dimension)).astype(np.float32)
|
|
vector_ids = [f"clear_test_{i}" for i in range(5)]
|
|
|
|
self.manager.add_vectors_batch(vectors, vector_ids)
|
|
|
|
# 验证向量已添加
|
|
self.assertEqual(self.manager.index.ntotal, 5)
|
|
self.assertEqual(len(self.manager.reverse_id_mapping), 5)
|
|
|
|
# 清空索引
|
|
result = self.manager.clear_index()
|
|
self.assertTrue(result)
|
|
|
|
# 验证索引已清空
|
|
self.assertEqual(self.manager.index.ntotal, 0)
|
|
self.assertEqual(len(self.manager.reverse_id_mapping), 0)
|
|
self.assertEqual(len(self.manager.metadata), 0)
|
|
self.assertEqual(self.manager.next_faiss_id, 0)
|
|
|
|
def test_get_index_info(self):
|
|
"""测试获取索引信息"""
|
|
# 初始状态
|
|
info = self.manager.get_index_info()
|
|
self.assertEqual(info['total_vectors'], 0)
|
|
self.assertEqual(info['dimension'], self.dimension)
|
|
self.assertEqual(info['active_vectors'], 0)
|
|
|
|
# 添加向量后
|
|
vectors = np.random.random((3, self.dimension)).astype(np.float32)
|
|
vector_ids = [f"info_test_{i}" for i in range(3)]
|
|
|
|
self.manager.add_vectors_batch(vectors, vector_ids)
|
|
|
|
info = self.manager.get_index_info()
|
|
self.assertEqual(info['total_vectors'], 3)
|
|
self.assertEqual(info['active_vectors'], 3)
|
|
|
|
# 删除一个向量后
|
|
self.manager.delete_vector(vector_ids[0])
|
|
|
|
info = self.manager.get_index_info()
|
|
self.assertEqual(info['total_vectors'], 3) # FAISS索引中的总数不变
|
|
self.assertEqual(info['active_vectors'], 2) # 活跃向量数减少
|
|
|
|
def test_different_index_types(self):
|
|
"""测试不同的索引类型"""
|
|
index_types = ["IndexFlatIP", "IndexFlatL2"]
|
|
|
|
for index_type in index_types:
|
|
with self.subTest(index_type=index_type):
|
|
manager = FAISSManager(
|
|
dimension=self.dimension,
|
|
index_type=index_type
|
|
)
|
|
|
|
# 添加向量
|
|
vector = np.random.random(self.dimension).astype(np.float32)
|
|
result = manager.add_vector(vector, f"test_{index_type}")
|
|
self.assertTrue(result)
|
|
|
|
# 搜索
|
|
query = np.random.random(self.dimension).astype(np.float32)
|
|
results = manager.search_similar(query, k=1)
|
|
self.assertEqual(len(results), 1)
|
|
|
|
def test_vector_normalization(self):
|
|
"""测试向量标准化"""
|
|
# 测试标准化开启
|
|
manager_norm = FAISSManager(
|
|
dimension=self.dimension,
|
|
normalize_vectors=True
|
|
)
|
|
|
|
# 测试标准化关闭
|
|
manager_no_norm = FAISSManager(
|
|
dimension=self.dimension,
|
|
normalize_vectors=False
|
|
)
|
|
|
|
# 添加相同的向量
|
|
vector = np.random.random(self.dimension).astype(np.float32) * 10 # 放大向量
|
|
|
|
manager_norm.add_vector(vector, "norm_test")
|
|
manager_no_norm.add_vector(vector, "no_norm_test")
|
|
|
|
# 两个管理器都应该成功添加
|
|
self.assertEqual(manager_norm.index.ntotal, 1)
|
|
self.assertEqual(manager_no_norm.index.ntotal, 1)
|
|
|
|
|
|
class TestFAISSManagerIntegration(unittest.TestCase):
|
|
"""FAISS管理器集成测试"""
|
|
|
|
def test_food_classification_scenario(self):
|
|
"""测试食物分类场景"""
|
|
# 创建管理器
|
|
manager = FAISSManager(dimension=128, index_type="IndexFlatIP")
|
|
|
|
# 模拟食物向量数据
|
|
food_data = [
|
|
{"id": "sichuan_001", "name": "回锅肉", "category": "川菜", "spicy": 3},
|
|
{"id": "sichuan_002", "name": "麻辣小面", "category": "川菜", "spicy": 4},
|
|
{"id": "sichuan_003", "name": "宫保鸡丁", "category": "川菜", "spicy": 3},
|
|
{"id": "home_001", "name": "西红柿鸡蛋", "category": "家常菜", "spicy": 0},
|
|
{"id": "home_002", "name": "炒细面", "category": "家常菜", "spicy": 1},
|
|
]
|
|
|
|
# 生成模拟向量（实际应用中这些是通过embedding模型生成的）
|
|
vectors = []
|
|
for i, food in enumerate(food_data):
|
|
# 为川菜生成相似的向量，为家常菜生成另一类相似的向量
|
|
if food["category"] == "川菜":
|
|
base_vector = np.array([1.0] * 64 + [0.0] * 64)
|
|
else:
|
|
base_vector = np.array([0.0] * 64 + [1.0] * 64)
|
|
|
|
# 添加随机噪声
|
|
noise = np.random.normal(0, 0.1, 128)
|
|
vector = (base_vector + noise).astype(np.float32)
|
|
vectors.append(vector)
|
|
|
|
# 批量添加向量
|
|
vector_ids = [food["id"] for food in food_data]
|
|
metadata_list = [
|
|
{"name": food["name"], "category": food["category"], "spicy": food["spicy"]}
|
|
for food in food_data
|
|
]
|
|
|
|
results = manager.add_vectors_batch(vectors, vector_ids, metadata_list)
|
|
self.assertTrue(all(results))
|
|
|
|
# 测试川菜查询
|
|
sichuan_query = np.array([1.0] * 64 + [0.0] * 64).astype(np.float32)
|
|
sichuan_results = manager.search_similar(sichuan_query, k=3)
|
|
|
|
# 验证川菜结果
|
|
self.assertGreater(len(sichuan_results), 0)
|
|
for vector_id, similarity, metadata in sichuan_results:
|
|
self.assertEqual(metadata["category"], "川菜")
|
|
|
|
# 测试家常菜查询
|
|
home_query = np.array([0.0] * 64 + [1.0] * 64).astype(np.float32)
|
|
home_results = manager.search_similar(home_query, k=3)
|
|
|
|
# 验证家常菜结果
|
|
self.assertGreater(len(home_results), 0)
|
|
for vector_id, similarity, metadata in home_results:
|
|
self.assertEqual(metadata["category"], "家常菜")
|
|
|
|
# 测试保存和加载
|
|
temp_dir = tempfile.mkdtemp()
|
|
try:
|
|
save_path = os.path.join(temp_dir, "food_index")
|
|
self.assertTrue(manager.save_index(save_path))
|
|
|
|
# 加载到新管理器
|
|
new_manager = FAISSManager(dimension=128)
|
|
self.assertTrue(new_manager.load_index(save_path))
|
|
|
|
# 验证加载后的搜索功能
|
|
loaded_results = new_manager.search_similar(sichuan_query, k=2)
|
|
self.assertGreater(len(loaded_results), 0)
|
|
|
|
finally:
|
|
shutil.rmtree(temp_dir)
|
|
|
|
|
|
def run_tests():
|
|
"""运行所有测试"""
|
|
# 创建测试套件
|
|
test_suite = unittest.TestSuite()
|
|
|
|
# 添加基本功能测试
|
|
test_suite.addTest(unittest.makeSuite(TestFAISSManager))
|
|
|
|
# 添加集成测试
|
|
test_suite.addTest(unittest.makeSuite(TestFAISSManagerIntegration))
|
|
|
|
# 运行测试
|
|
runner = unittest.TextTestRunner(verbosity=2)
|
|
result = runner.run(test_suite)
|
|
|
|
return result.wasSuccessful()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
print("开始运行FAISS管理器测试...")
|
|
success = run_tests()
|
|
|
|
if success:
|
|
print("\n✅ 所有测试通过！")
|
|
else:
|
|
print("\n❌ 部分测试失败，请检查代码。")
|
|
sys.exit(1) |