"""嵌入模型测试.""" import pytest from src.core.config import EmbedConfig from src.core.embedder import LocalEmbedder, APIEmbedder, create_embedder, batch_embed class TestLocalEmbedder: """本地嵌入器测试.""" @pytest.fixture def local_config(self): return EmbedConfig(mode="local") def test_create_local_embedder(self, local_config): """创建本地嵌入器,验证维度正确.""" embedder = create_embedder(local_config) assert embedder.dimension > 0 assert isinstance(embedder.dimension, int) assert isinstance(embedder, LocalEmbedder) def test_embed_single_text(self, local_config): """嵌入单条文本返回正确维度向量.""" embedder = create_embedder(local_config) result = embedder.embed(["你好世界"]) assert len(result) == 1 assert len(result[0]) == embedder.dimension assert all(isinstance(v, float) for v in result[0]) def test_embed_multiple_texts(self, local_config): """嵌入多条文本返回对应数量的向量.""" embedder = create_embedder(local_config) texts = ["第一段文本", "第二段文本", "第三段文本"] result = embedder.embed(texts) assert len(result) == 3 for vec in result: assert len(vec) == embedder.dimension def test_embed_empty_list_raises(self, local_config): """空列表应抛出异常.""" embedder = create_embedder(local_config) with pytest.raises(ValueError): embedder.embed([]) class TestAPIEmbedder: """API 嵌入器测试.""" def test_api_embedder_init(self): """API 嵌入器初始化.""" cfg = EmbedConfig(mode="api", api_base="https://api.test.com", api_key="sk-test") emb = APIEmbedder(cfg) assert emb.dimension == 1536 def test_api_embedder_empty_raises(self): """API 嵌入器空列表抛异常.""" cfg = EmbedConfig(mode="api") emb = APIEmbedder(cfg) with pytest.raises(ValueError): emb.embed([]) class TestBatchEmbed: """分批嵌入测试.""" def test_batch_embed(self): """分批嵌入返回正确数量.""" embedder = create_embedder(EmbedConfig(mode="local")) texts = ["测试文本"] * 70 # 超过 batch_size 32 results = batch_embed(embedder, texts, batch_size=32) assert len(results) == 70 assert len(results[0]) == embedder.dimension