250 lines
8.7 KiB
Python
250 lines
8.7 KiB
Python
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",
|
|
)
|