import asyncio from collections.abc import AsyncIterator import pytest import pytest_asyncio from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine from deerflow.persistence.assistant_knowledge import AssistantKnowledgeRepository from deerflow.persistence.base import Base class SnapshotEmbeddingClient: dimensions = 3 fingerprint = "snapshot-test-model" async def embed_query(self, _query: str) -> list[float]: return [1.0, 0.0, 0.0] @pytest_asyncio.fixture async def repository(tmp_path) -> AsyncIterator[AssistantKnowledgeRepository]: # A file-backed database gives the atomic-switch test a second connection # that can read the last committed snapshot while the replacement # transaction is deliberately paused before commit. database_path = (tmp_path / "assistant-snapshot.sqlite3").as_posix() engine = create_async_engine(f"sqlite+aiosqlite:///{database_path}") async with engine.begin() as connection: await connection.run_sync(Base.metadata.create_all) yield AssistantKnowledgeRepository(async_sessionmaker(engine, expire_on_commit=False)) await engine.dispose() def _package(*, revision: int, marker: str) -> dict: return { "source": { "provider": "weknora", "type": "weknora", "knowledge_base_id": "stable-kb", "knowledge_base_name": "稳定普通知识库", "index_revision": revision, }, "wiki_pages": [ { "slug": "wiki/stable", "title": "稳定页面", "page_type": "article", "content": f"# 稳定页面\n\n{marker}", } ], "vectors": [ { "wiki_slug": "wiki/stable", "section_index": 0, "section_content": marker, "embedding_model": "snapshot-test-model", "embedding_fingerprint": SnapshotEmbeddingClient.fingerprint, "embedding_dimensions": SnapshotEmbeddingClient.dimensions, "vector": [1.0, 0.0, 0.0], } ], } async def _seed_snapshot( repository: AssistantKnowledgeRepository, ) -> tuple[dict, dict, dict]: base = await repository.create_base(created_by="admin", name="助手快照测试库") job = await repository.import_package( package=_package(revision=1, marker="OLD-SNAPSHOT-MARKER-991"), source_type="weknora", source_key="stable-kb", source_name="稳定普通知识库", trigger="initial_import", created_by="admin", mapping_id="stable-kb", base_id=base["id"], mirror_to_global=False, ) source = (await repository.list_sources(base_id=base["id"]))[0] assert source["last_import_job_id"] == job["id"] return base, source, job async def _assert_vector_snapshot( repository: AssistantKnowledgeRepository, *, base_id: str, marker: str, ) -> None: results = await repository.search_pages( base_ids=[base_id], query="完全没有字面重合的向量查询", embedding_client=SnapshotEmbeddingClient(), ) assert results assert results[0]["matched_section"]["content"] == marker async def _queue_reimport( repository: AssistantKnowledgeRepository, *, base_id: str, revision: int, ) -> dict: return await repository.create_queued_import_job( source_type="weknora", source_key="stable-kb", source_name="稳定普通知识库（待发布）", trigger="source_reimport", created_by="admin", metadata={ "provider": "weknora", "knowledge_base_id": "stable-kb", "knowledge_base_name": "稳定普通知识库（待发布）", "index_revision": revision, }, mapping_id="stable-kb", base_id=base_id, ) @pytest.mark.asyncio async def test_queued_and_running_reimport_keep_last_successful_snapshot( repository: AssistantKnowledgeRepository, ) -> None: base, original_source, original_job = await _seed_snapshot(repository) queued_job = await _queue_reimport(repository, base_id=base["id"], revision=2) queued_source = await repository.get_source(original_source["id"]) assert queued_job["status"] == "queued" assert queued_source is not None assert queued_source["last_import_job_id"] == original_job["id"] assert queued_source["metadata"]["index_revision"] == 1 assert queued_source["source_name"] == "稳定普通知识库" await _assert_vector_snapshot( repository, base_id=base["id"], marker="OLD-SNAPSHOT-MARKER-991", ) await repository.update_import_job( queued_job["id"], status="running", phase="exporting", ) running_source = await repository.get_source(original_source["id"]) assert running_source is not None assert running_source["last_import_job_id"] == original_job["id"] await _assert_vector_snapshot( repository, base_id=base["id"], marker="OLD-SNAPSHOT-MARKER-991", ) @pytest.mark.asyncio async def test_failed_reimport_rolls_back_and_keeps_old_vector_snapshot( repository: AssistantKnowledgeRepository, ) -> None: base, original_source, original_job = await _seed_snapshot(repository) queued_job = await _queue_reimport(repository, base_id=base["id"], revision=2) invalid_package = _package(revision=2, marker="BROKEN-NEW-SNAPSHOT") invalid_package["vectors"][0]["wiki_slug"] = "wiki/missing" with pytest.raises(ValueError, match="向量引用不存在"): await repository.import_package_for_job( queued_job["id"], package=invalid_package, created_by="admin", mirror_to_global=False, ) await repository.update_import_job( queued_job["id"], status="failed", phase="failed", error="测试导入失败", completed=True, ) failed_source = await repository.get_source(original_source["id"]) assert failed_source is not None assert failed_source["last_import_job_id"] == original_job["id"] assert failed_source["metadata"]["index_revision"] == 1 await _assert_vector_snapshot( repository, base_id=base["id"], marker="OLD-SNAPSHOT-MARKER-991", ) pages = await repository.list_pages(base_id=base["id"]) assert pages[0]["content_markdown"].endswith("OLD-SNAPSHOT-MARKER-991") @pytest.mark.asyncio async def test_successful_reimport_switches_snapshot_only_at_commit( repository: AssistantKnowledgeRepository, monkeypatch: pytest.MonkeyPatch, ) -> None: base, original_source, original_job = await _seed_snapshot(repository) queued_job = await _queue_reimport(repository, base_id=base["id"], revision=2) snapshot_written = asyncio.Event() allow_commit = asyncio.Event() original_retire_orphans = repository._retire_orphans async def pause_before_snapshot_switch(session, base_id: str) -> None: await original_retire_orphans(session, base_id) snapshot_written.set() await allow_commit.wait() monkeypatch.setattr(repository, "_retire_orphans", pause_before_snapshot_switch) import_task = asyncio.create_task( repository.import_package_for_job( queued_job["id"], package=_package(revision=2, marker="NEW-SNAPSHOT-MARKER-772"), created_by="admin", mirror_to_global=False, ) ) await asyncio.wait_for(snapshot_written.wait(), timeout=5) try: # The replacement rows exist only in the writer's uncommitted # transaction. Other readers must still resolve the old job pointer and # the old vector snapshot until the transaction commits. mid_import_source = await repository.get_source(original_source["id"]) assert mid_import_source is not None assert mid_import_source["last_import_job_id"] == original_job["id"] assert mid_import_source["metadata"]["index_revision"] == 1 await _assert_vector_snapshot( repository, base_id=base["id"], marker="OLD-SNAPSHOT-MARKER-991", ) finally: allow_commit.set() completed_job = await asyncio.wait_for(import_task, timeout=5) published_source = await repository.get_source(original_source["id"]) assert completed_job["status"] == "completed" assert published_source is not None assert published_source["last_import_job_id"] == queued_job["id"] assert published_source["metadata"]["index_revision"] == 2 assert published_source["source_name"] == "稳定普通知识库" await _assert_vector_snapshot( repository, base_id=base["id"], marker="NEW-SNAPSHOT-MARKER-772", )