mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-15 17:30:13 +08:00
296fcfb54a
* fix: roll back document rows when FAISS insert fails Validate embedding shape/dimension before local writes. If DocumentStorage commits but FAISS write fails, best-effort delete the partial FAISS ids and document rows so the two stores do not diverge. * fix: clean up partial KB upload state before metadata commit Track metadata_committed right after commit succeeds. On failure before metadata commit, roll back chunks/vectors first, then residual KB rows and media files. After metadata is committed, keep the document even if stats or refresh fails. * test: cover knowledge base upload atomicity and rollback Add unit tests for insert_batch FAISS failure rollback, dimension validation before document writes, upload cleanup order, no-rollback after metadata commit, and real DocumentStorage+FAISS no-orphan cases. * fix: log media directory cleanup failures after upload rollback Best-effort rmdir of the per-document media directory should still emit a warning when removal fails so incomplete cleanup is diagnosable.
575 lines
19 KiB
Python
575 lines
19 KiB
Python
"""
|
|
Unit tests for knowledge base upload atomicity / compensating rollback.
|
|
|
|
Covers:
|
|
1. insert_batch rolls back document rows when FAISS insert fails
|
|
2. upload_document cleans up chunks/vectors/media on storage failure
|
|
3. upload_document cleans up after metadata failure (post insert_batch)
|
|
4. upload_document does NOT roll back after metadata is committed
|
|
5. Real DocumentStorage + EmbeddingStorage leave no orphans on FAISS failure
|
|
6. Dimension validation failures never write to DocumentStorage
|
|
|
|
These tests use lazy imports and a ProviderManager stub to avoid circular
|
|
import issues in the astrbot core module chain (same pattern as
|
|
test_kb_manager_resilience.py).
|
|
"""
|
|
|
|
import sys
|
|
import types
|
|
from contextlib import asynccontextmanager
|
|
from pathlib import Path
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from astrbot.core.db.vec_db.faiss_impl.vec_db import FaissVecDB
|
|
from astrbot.core.exceptions import KnowledgeBaseUploadError
|
|
from astrbot.core.knowledge_base.models import KnowledgeBase
|
|
|
|
|
|
@pytest.fixture
|
|
def stub_provider_manager_module():
|
|
"""Stub provider manager module to avoid circular imports in unit tests."""
|
|
original_module = sys.modules.get("astrbot.core.provider.manager")
|
|
stub_module = types.ModuleType("astrbot.core.provider.manager")
|
|
|
|
class ProviderManager: ...
|
|
|
|
setattr(stub_module, "ProviderManager", ProviderManager)
|
|
sys.modules["astrbot.core.provider.manager"] = stub_module
|
|
|
|
# Drop already-imported modules that transitively need ProviderManager so
|
|
# they re-import against the stub.
|
|
to_drop = [
|
|
name
|
|
for name in list(sys.modules)
|
|
if name.startswith("astrbot.core.knowledge_base.kb_helper")
|
|
or name.startswith("astrbot.core.knowledge_base.kb_mgr")
|
|
]
|
|
for name in to_drop:
|
|
sys.modules.pop(name, None)
|
|
|
|
try:
|
|
yield
|
|
finally:
|
|
if original_module is not None:
|
|
sys.modules["astrbot.core.provider.manager"] = original_module
|
|
else:
|
|
sys.modules.pop("astrbot.core.provider.manager", None)
|
|
|
|
|
|
def _make_vec_db() -> FaissVecDB:
|
|
vec_db = FaissVecDB.__new__(FaissVecDB)
|
|
vec_db.embedding_provider = AsyncMock()
|
|
vec_db.document_storage = AsyncMock()
|
|
vec_db.embedding_storage = AsyncMock()
|
|
vec_db.embedding_storage.dimension = 2
|
|
return vec_db
|
|
|
|
|
|
def _import_kb_helper():
|
|
from astrbot.core.knowledge_base.kb_helper import KBHelper
|
|
|
|
return KBHelper
|
|
|
|
|
|
def _failing_get_db():
|
|
@asynccontextmanager
|
|
async def _cm():
|
|
raise RuntimeError("kb.db locked")
|
|
yield # pragma: no cover
|
|
|
|
return _cm
|
|
|
|
|
|
def _successful_get_db(session):
|
|
@asynccontextmanager
|
|
async def _cm():
|
|
yield session
|
|
|
|
return _cm
|
|
|
|
|
|
def _successful_begin():
|
|
@asynccontextmanager
|
|
async def _cm():
|
|
yield None
|
|
|
|
return _cm
|
|
|
|
|
|
def _session_with_begin(execute_side_effect=None):
|
|
"""Session mock that supports ``async with session.begin()``."""
|
|
session = MagicMock()
|
|
session.begin = _successful_begin()
|
|
session.add = MagicMock()
|
|
session.commit = AsyncMock()
|
|
session.refresh = AsyncMock()
|
|
session.execute = AsyncMock(side_effect=execute_side_effect)
|
|
return session
|
|
|
|
|
|
async def _make_real_vec_db(tmp_path: Path, dim: int = 4) -> FaissVecDB:
|
|
"""Build a FaissVecDB backed by real DocumentStorage + EmbeddingStorage."""
|
|
doc_path = str(tmp_path / "doc.db")
|
|
index_path = str(tmp_path / "index.faiss")
|
|
|
|
embedding_provider = MagicMock()
|
|
# get_dim is sync in EmbeddingProvider; must return a plain int for FAISS.
|
|
embedding_provider.get_dim = MagicMock(return_value=dim)
|
|
embedding_provider.get_embeddings_batch = AsyncMock()
|
|
embedding_provider.get_embedding = AsyncMock()
|
|
|
|
vec_db = FaissVecDB(
|
|
doc_store_path=doc_path,
|
|
index_store_path=index_path,
|
|
embedding_provider=embedding_provider,
|
|
)
|
|
await vec_db.initialize()
|
|
return vec_db
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_insert_batch_rolls_back_documents_when_faiss_fails() -> None:
|
|
"""FAISS write failure should delete rows already committed to DocumentStorage."""
|
|
vec_db = _make_vec_db()
|
|
vec_db.embedding_provider.get_embeddings_batch.return_value = [
|
|
[0.1, 0.2],
|
|
[0.3, 0.4],
|
|
]
|
|
vec_db.document_storage.insert_documents_batch.return_value = [11, 12]
|
|
vec_db.embedding_storage.insert_batch.side_effect = RuntimeError(
|
|
"faiss write failed",
|
|
)
|
|
|
|
with pytest.raises(RuntimeError, match="faiss write failed"):
|
|
await FaissVecDB.insert_batch(
|
|
vec_db,
|
|
contents=["chunk-1", "chunk-2"],
|
|
metadatas=[{"kb_doc_id": "doc-1"}, {"kb_doc_id": "doc-1"}],
|
|
ids=["c1", "c2"],
|
|
)
|
|
|
|
vec_db.embedding_storage.delete.assert_awaited_once_with([11, 12])
|
|
assert vec_db.document_storage.delete_document_by_doc_id.await_count == 2
|
|
vec_db.document_storage.delete_document_by_doc_id.assert_any_await("c1")
|
|
vec_db.document_storage.delete_document_by_doc_id.assert_any_await("c2")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_insert_batch_rolls_back_on_int_id_count_mismatch() -> None:
|
|
"""Mismatched int_id count after document insert should roll back those docs."""
|
|
vec_db = _make_vec_db()
|
|
vec_db.embedding_provider.get_embeddings_batch.return_value = [
|
|
[0.1, 0.2],
|
|
[0.3, 0.4],
|
|
]
|
|
vec_db.document_storage.insert_documents_batch.return_value = [11] # mismatch
|
|
|
|
with pytest.raises(KnowledgeBaseUploadError) as exc_info:
|
|
await FaissVecDB.insert_batch(
|
|
vec_db,
|
|
contents=["chunk-1", "chunk-2"],
|
|
metadatas=[{}, {}],
|
|
ids=["c1", "c2"],
|
|
)
|
|
|
|
assert "内部 ID 数量" in str(exc_info.value)
|
|
vec_db.embedding_storage.insert_batch.assert_not_awaited()
|
|
vec_db.embedding_storage.delete.assert_awaited_once_with([11])
|
|
assert vec_db.document_storage.delete_document_by_doc_id.await_count == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_insert_batch_dimension_mismatch_does_not_write_documents() -> None:
|
|
"""Dimension validation must fail before DocumentStorage is written."""
|
|
vec_db = _make_vec_db()
|
|
vec_db.embedding_storage.dimension = 4
|
|
vec_db.embedding_provider.get_embeddings_batch.return_value = [
|
|
[0.1, 0.2], # wrong dim
|
|
[0.3, 0.4],
|
|
]
|
|
|
|
with pytest.raises(KnowledgeBaseUploadError) as exc_info:
|
|
await FaissVecDB.insert_batch(
|
|
vec_db,
|
|
contents=["chunk-1", "chunk-2"],
|
|
metadatas=[{}, {}],
|
|
ids=["c1", "c2"],
|
|
)
|
|
|
|
assert "维度" in str(exc_info.value)
|
|
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_real_storage_rolls_back_on_faiss_failure(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""Real DocumentStorage must not keep orphan chunks when FAISS write fails."""
|
|
vec_db = await _make_real_vec_db(tmp_path, dim=4)
|
|
vec_db.embedding_provider.get_embeddings_batch.return_value = [
|
|
[0.1, 0.2, 0.3, 0.4],
|
|
[0.5, 0.6, 0.7, 0.8],
|
|
]
|
|
|
|
original_insert = vec_db.embedding_storage.insert_batch
|
|
|
|
async def boom(vectors, ids):
|
|
# Simulate failure after in-memory add would happen: force raise before
|
|
# success so DocumentStorage rows must be cleaned up.
|
|
raise RuntimeError("simulated faiss disk write failure")
|
|
|
|
vec_db.embedding_storage.insert_batch = boom # type: ignore[method-assign]
|
|
|
|
kb_doc_id = "kb-doc-real-1"
|
|
with pytest.raises(RuntimeError, match="simulated faiss"):
|
|
await vec_db.insert_batch(
|
|
contents=["alpha chunk", "beta chunk"],
|
|
metadatas=[
|
|
{"kb_id": "kb-1", "kb_doc_id": kb_doc_id, "chunk_index": 0},
|
|
{"kb_id": "kb-1", "kb_doc_id": kb_doc_id, "chunk_index": 1},
|
|
],
|
|
ids=["chunk-a", "chunk-b"],
|
|
)
|
|
|
|
remaining = await vec_db.document_storage.get_documents(
|
|
metadata_filters={"kb_doc_id": kb_doc_id},
|
|
offset=None,
|
|
limit=None,
|
|
)
|
|
assert remaining == []
|
|
assert await vec_db.count_documents(metadata_filter={"kb_doc_id": kb_doc_id}) == 0
|
|
|
|
# Restore for clean close; index should still be empty / consistent.
|
|
vec_db.embedding_storage.insert_batch = original_insert # type: ignore[method-assign]
|
|
await vec_db.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_insert_batch_real_storage_dimension_mismatch_leaves_no_docs(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""Real storage: wrong embedding dim must not leave any document rows."""
|
|
vec_db = await _make_real_vec_db(tmp_path, dim=4)
|
|
vec_db.embedding_provider.get_embeddings_batch.return_value = [
|
|
[0.1, 0.2], # dim 2 != 4
|
|
[0.3, 0.4],
|
|
]
|
|
|
|
with pytest.raises(KnowledgeBaseUploadError):
|
|
await vec_db.insert_batch(
|
|
contents=["alpha", "beta"],
|
|
metadatas=[
|
|
{"kb_id": "kb-1", "kb_doc_id": "doc-dim", "chunk_index": 0},
|
|
{"kb_id": "kb-1", "kb_doc_id": "doc-dim", "chunk_index": 1},
|
|
],
|
|
ids=["c1", "c2"],
|
|
)
|
|
|
|
remaining = await vec_db.document_storage.get_documents(
|
|
metadata_filters={"kb_doc_id": "doc-dim"},
|
|
offset=None,
|
|
limit=None,
|
|
)
|
|
assert remaining == []
|
|
await vec_db.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_upload_document_cleans_up_on_storage_failure(
|
|
tmp_path: Path,
|
|
stub_provider_manager_module,
|
|
) -> None:
|
|
"""Storage failure should clean media and request chunk/vector rollback."""
|
|
KBHelper = _import_kb_helper()
|
|
|
|
helper = KBHelper.__new__(KBHelper)
|
|
helper.kb = KnowledgeBase(
|
|
kb_name="Test KB",
|
|
description="",
|
|
embedding_provider_id="emb",
|
|
)
|
|
helper.kb_db = MagicMock()
|
|
helper.vec_db = AsyncMock()
|
|
helper.kb_medias_dir = tmp_path / "medias"
|
|
helper.kb_medias_dir.mkdir()
|
|
helper.chunker = AsyncMock()
|
|
helper.chunker.chunk = AsyncMock(return_value=["hello world"])
|
|
|
|
media_file = helper.kb_medias_dir / "will-be-set" / "img.png"
|
|
|
|
async def fake_save_media(**kwargs):
|
|
nonlocal media_file
|
|
doc_id = kwargs["doc_id"]
|
|
media_dir = helper.kb_medias_dir / doc_id
|
|
media_dir.mkdir(parents=True, exist_ok=True)
|
|
media_file = media_dir / "img.png"
|
|
media_file.write_bytes(b"fake-image")
|
|
media = MagicMock()
|
|
media.file_path = str(media_file)
|
|
return media
|
|
|
|
helper._save_media = AsyncMock(side_effect=fake_save_media)
|
|
helper.vec_db.insert_batch.side_effect = RuntimeError("embedding provider down")
|
|
helper.vec_db.delete_documents = AsyncMock()
|
|
helper.kb_db.get_db = _successful_get_db(_session_with_begin())
|
|
|
|
parse_result = MagicMock()
|
|
parse_result.text = "hello world"
|
|
parse_result.media = [
|
|
MagicMock(
|
|
media_type="image",
|
|
file_name="img.png",
|
|
content=b"fake-image",
|
|
mime_type="image/png",
|
|
),
|
|
]
|
|
|
|
with (
|
|
patch(
|
|
"astrbot.core.knowledge_base.kb_helper.select_parser",
|
|
new=AsyncMock(
|
|
return_value=MagicMock(parse=AsyncMock(return_value=parse_result)),
|
|
),
|
|
),
|
|
patch.object(helper, "_ensure_vec_db", new=AsyncMock()),
|
|
pytest.raises(KnowledgeBaseUploadError) as exc_info,
|
|
):
|
|
await helper.upload_document(
|
|
file_name="demo.txt",
|
|
file_content=b"hello world",
|
|
file_type="txt",
|
|
)
|
|
|
|
assert exc_info.value.stage == "storage"
|
|
assert "cause" in exc_info.value.details
|
|
helper.vec_db.delete_documents.assert_awaited()
|
|
assert not media_file.exists()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_upload_document_cleans_up_on_metadata_failure(
|
|
stub_provider_manager_module,
|
|
) -> None:
|
|
"""Metadata commit failure after insert_batch should delete written chunks."""
|
|
KBHelper = _import_kb_helper()
|
|
|
|
helper = KBHelper.__new__(KBHelper)
|
|
helper.kb = KnowledgeBase(
|
|
kb_name="Test KB",
|
|
description="",
|
|
embedding_provider_id="emb",
|
|
)
|
|
helper.kb_db = MagicMock()
|
|
helper.vec_db = AsyncMock()
|
|
helper.kb_medias_dir = Path("/tmp/kb-medias-unused")
|
|
helper.chunker = AsyncMock()
|
|
|
|
helper.vec_db.insert_batch = AsyncMock()
|
|
helper.vec_db.delete_documents = AsyncMock()
|
|
helper.kb_db.get_db = _failing_get_db()
|
|
|
|
with (
|
|
patch.object(helper, "_ensure_vec_db", new=AsyncMock()),
|
|
pytest.raises(KnowledgeBaseUploadError) as exc_info,
|
|
):
|
|
await helper.upload_document(
|
|
file_name="demo.txt",
|
|
file_content=None,
|
|
file_type="txt",
|
|
pre_chunked_text=["chunk a", "chunk b"],
|
|
)
|
|
|
|
assert exc_info.value.stage == "metadata"
|
|
helper.vec_db.insert_batch.assert_awaited_once()
|
|
helper.vec_db.delete_documents.assert_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_upload_document_skips_rollback_after_metadata_commit(
|
|
stub_provider_manager_module,
|
|
) -> None:
|
|
"""Stats refresh failure after metadata commit must not roll back the doc."""
|
|
KBHelper = _import_kb_helper()
|
|
|
|
helper = KBHelper.__new__(KBHelper)
|
|
helper.kb = KnowledgeBase(
|
|
kb_name="Test KB",
|
|
description="",
|
|
embedding_provider_id="emb",
|
|
)
|
|
helper.kb_db = MagicMock()
|
|
helper.vec_db = AsyncMock()
|
|
helper.kb_medias_dir = Path("/tmp/kb-medias-unused")
|
|
helper.chunker = AsyncMock()
|
|
helper.vec_db.insert_batch = AsyncMock()
|
|
helper.vec_db.delete_documents = AsyncMock()
|
|
helper.kb_db.update_kb_stats = AsyncMock(side_effect=RuntimeError("stats fail"))
|
|
helper.refresh_kb = AsyncMock()
|
|
helper.refresh_document = AsyncMock()
|
|
|
|
session = _session_with_begin()
|
|
helper.kb_db.get_db = _successful_get_db(session)
|
|
|
|
with (
|
|
patch.object(helper, "_ensure_vec_db", new=AsyncMock()),
|
|
pytest.raises(KnowledgeBaseUploadError) as exc_info,
|
|
):
|
|
await helper.upload_document(
|
|
file_name="demo.txt",
|
|
file_content=None,
|
|
file_type="txt",
|
|
pre_chunked_text=["chunk a"],
|
|
)
|
|
|
|
assert exc_info.value.stage == "metadata"
|
|
assert "统计信息刷新失败" in str(exc_info.value)
|
|
helper.vec_db.delete_documents.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_upload_document_skips_rollback_when_refresh_fails_after_commit(
|
|
stub_provider_manager_module,
|
|
) -> None:
|
|
"""If commit succeeds but session.refresh fails, do not roll back."""
|
|
KBHelper = _import_kb_helper()
|
|
|
|
helper = KBHelper.__new__(KBHelper)
|
|
helper.kb = KnowledgeBase(
|
|
kb_name="Test KB",
|
|
description="",
|
|
embedding_provider_id="emb",
|
|
)
|
|
helper.kb_db = MagicMock()
|
|
helper.vec_db = AsyncMock()
|
|
helper.kb_medias_dir = Path("/tmp/kb-medias-unused")
|
|
helper.chunker = AsyncMock()
|
|
helper.vec_db.insert_batch = AsyncMock()
|
|
helper.vec_db.delete_documents = AsyncMock()
|
|
|
|
session = _session_with_begin()
|
|
session.refresh = AsyncMock(side_effect=RuntimeError("refresh failed"))
|
|
helper.kb_db.get_db = _successful_get_db(session)
|
|
|
|
with (
|
|
patch.object(helper, "_ensure_vec_db", new=AsyncMock()),
|
|
pytest.raises(KnowledgeBaseUploadError) as exc_info,
|
|
):
|
|
await helper.upload_document(
|
|
file_name="demo.txt",
|
|
file_content=None,
|
|
file_type="txt",
|
|
pre_chunked_text=["chunk a"],
|
|
)
|
|
|
|
assert exc_info.value.stage == "metadata"
|
|
assert "文档记录刷新失败" in str(exc_info.value)
|
|
helper.vec_db.delete_documents.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cleanup_failed_upload_deletes_vectors_before_metadata(
|
|
tmp_path: Path,
|
|
stub_provider_manager_module,
|
|
) -> None:
|
|
"""Rollback should delete vectors first, then metadata rows, then media."""
|
|
KBHelper = _import_kb_helper()
|
|
|
|
helper = KBHelper.__new__(KBHelper)
|
|
helper.vec_db = AsyncMock()
|
|
helper.kb_db = MagicMock()
|
|
helper.kb_medias_dir = tmp_path / "medias"
|
|
helper.kb_medias_dir.mkdir()
|
|
|
|
call_order: list[str] = []
|
|
|
|
async def track_delete_documents(**kwargs):
|
|
call_order.append("vectors")
|
|
|
|
helper.vec_db.delete_documents = AsyncMock(side_effect=track_delete_documents)
|
|
|
|
async def track_execute(*args, **kwargs):
|
|
if "vectors" in call_order and "metadata" not in call_order:
|
|
call_order.append("metadata")
|
|
return None
|
|
|
|
helper.kb_db.get_db = _successful_get_db(
|
|
_session_with_begin(execute_side_effect=track_execute),
|
|
)
|
|
|
|
media = helper.kb_medias_dir / "doc-x" / "a.png"
|
|
media.parent.mkdir()
|
|
media.write_bytes(b"x")
|
|
|
|
await helper._cleanup_failed_upload(doc_id="doc-x", media_paths=[media])
|
|
|
|
assert call_order[0] == "vectors"
|
|
assert "metadata" in call_order
|
|
assert not media.exists()
|
|
helper.vec_db.delete_documents.assert_awaited_once_with(
|
|
metadata_filters={"kb_doc_id": "doc-x"},
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cleanup_failed_upload_is_best_effort(
|
|
tmp_path: Path,
|
|
stub_provider_manager_module,
|
|
) -> None:
|
|
"""Rollback path failures must not raise; media cleanup still runs."""
|
|
KBHelper = _import_kb_helper()
|
|
|
|
helper = KBHelper.__new__(KBHelper)
|
|
helper.kb_db = MagicMock()
|
|
helper.vec_db = AsyncMock()
|
|
helper.vec_db.delete_documents.side_effect = RuntimeError("vec delete failed")
|
|
helper.kb_db.get_db = _failing_get_db()
|
|
helper.kb_medias_dir = tmp_path / "medias"
|
|
helper.kb_medias_dir.mkdir()
|
|
|
|
media = helper.kb_medias_dir / "doc-x" / "a.png"
|
|
media.parent.mkdir()
|
|
media.write_bytes(b"x")
|
|
|
|
# Should not raise
|
|
await helper._cleanup_failed_upload(doc_id="doc-x", media_paths=[media])
|
|
assert not media.exists()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cleanup_failed_upload_real_vec_db_by_kb_doc_id(
|
|
tmp_path: Path,
|
|
stub_provider_manager_module,
|
|
) -> None:
|
|
"""Cleanup with a real vec_db removes chunks keyed by kb_doc_id."""
|
|
KBHelper = _import_kb_helper()
|
|
vec_db = await _make_real_vec_db(tmp_path, dim=4)
|
|
vec_db.embedding_provider.get_embeddings_batch.return_value = [
|
|
[0.1, 0.2, 0.3, 0.4],
|
|
[0.5, 0.6, 0.7, 0.8],
|
|
]
|
|
|
|
kb_doc_id = "upload-doc-cleanup"
|
|
await vec_db.insert_batch(
|
|
contents=["one", "two"],
|
|
metadatas=[
|
|
{"kb_id": "kb-1", "kb_doc_id": kb_doc_id, "chunk_index": 0},
|
|
{"kb_id": "kb-1", "kb_doc_id": kb_doc_id, "chunk_index": 1},
|
|
],
|
|
ids=["u1", "u2"],
|
|
)
|
|
assert await vec_db.count_documents(metadata_filter={"kb_doc_id": kb_doc_id}) == 2
|
|
|
|
helper = KBHelper.__new__(KBHelper)
|
|
helper.vec_db = vec_db
|
|
helper.kb_db = MagicMock()
|
|
helper.kb_medias_dir = tmp_path / "medias"
|
|
helper.kb_medias_dir.mkdir()
|
|
helper.kb_db.get_db = _successful_get_db(_session_with_begin())
|
|
|
|
await helper._cleanup_failed_upload(doc_id=kb_doc_id, media_paths=[])
|
|
assert await vec_db.count_documents(metadata_filter={"kb_doc_id": kb_doc_id}) == 0
|
|
await vec_db.close()
|