mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-16 01:40:15 +08:00
cc72c01c0e
* fix: improve knowledge base upload error messages * fix: deduplicate knowledge base upload logs * fix: handle type errors in kb embedding validation
47 lines
1.6 KiB
Python
47 lines
1.6 KiB
Python
from unittest.mock import AsyncMock
|
||
|
||
import pytest
|
||
|
||
from astrbot.core.db.vec_db.faiss_impl.vec_db import FaissVecDB
|
||
from astrbot.core.exceptions import KnowledgeBaseUploadError
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_insert_batch_skips_empty_contents() -> None:
|
||
vec_db = FaissVecDB.__new__(FaissVecDB)
|
||
vec_db.embedding_provider = AsyncMock()
|
||
vec_db.document_storage = AsyncMock()
|
||
vec_db.embedding_storage = AsyncMock()
|
||
|
||
result = await FaissVecDB.insert_batch(vec_db, [])
|
||
|
||
assert result == []
|
||
vec_db.embedding_provider.get_embeddings_batch.assert_not_awaited()
|
||
vec_db.document_storage.insert_documents_batch.assert_not_awaited()
|
||
vec_db.embedding_storage.insert_batch.assert_not_awaited()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_insert_batch_raises_friendly_error_for_embedding_count_mismatch() -> (
|
||
None
|
||
):
|
||
vec_db = FaissVecDB.__new__(FaissVecDB)
|
||
vec_db.embedding_provider = AsyncMock()
|
||
vec_db.embedding_provider.get_embeddings_batch.return_value = [[0.1, 0.2]]
|
||
vec_db.document_storage = AsyncMock()
|
||
vec_db.embedding_storage = AsyncMock()
|
||
vec_db.embedding_storage.dimension = 2
|
||
|
||
with pytest.raises(KnowledgeBaseUploadError) as exc_info:
|
||
await FaissVecDB.insert_batch(
|
||
vec_db,
|
||
contents=["chunk-1", "chunk-2"],
|
||
metadatas=[{}, {}],
|
||
ids=["doc-1", "doc-2"],
|
||
)
|
||
|
||
assert "向量化失败" in str(exc_info.value)
|
||
assert "期望 2，实际 1" in str(exc_info.value)
|
||
vec_db.document_storage.insert_documents_batch.assert_not_awaited()
|
||
vec_db.embedding_storage.insert_batch.assert_not_awaited()
|