zhifu-platform/tests/knowledge_recovery_acceptance.py

import asyncio
import json
import sys
from pathlib import Path
from uuid import uuid4

import pytest
from test_after_sales import scenario  # noqa: F401
from test_knowledge_memory import cleanup_vectors, document

from zhifu.modules.knowledge.config import KnowledgeSettings
from zhifu.modules.knowledge.recovery import rebuild, snapshot, verify_source
from zhifu.modules.knowledge.vector import connection, retrieval_client


async def test_rebuild_published_documents_and_current_permissions(scenario):  # noqa: F811
    target_name = "zhifu_recovery_" + uuid4().hex
    source = KnowledgeSettings()
    target = source.model_copy(update={"collection": target_name})
    cli_name = target_name + "_cli"
    concurrent_name = target_name + "_concurrent"
    call, app = scenario["call"], scenario["app"]
    original = app.state.knowledge
    paths = []
    try:
        paths.append(await document(scenario))
        for index in range(3):
            path = f"/admin/v1/stores/{scenario['store']}/knowledge/{uuid4()}"
            response = await call("PUT", path, "agent", {
                "expected_version": 0, "title": f"键盘恢复资料{index}",
                "text": f"键盘恢复资料{index}：检查接口连接，清洁前断开电源。",
            })
            assert response.status_code == 200, response.text
            paths.append(path)
        # 两份发布资料、一份草稿、一份撤回资料，来源均通过真实管理接口建立。
        for path in paths[:2] + paths[3:]:
            response = await call("POST", path + "/publish", "agent", {"expected_version": 1})
            assert response.status_code == 200, response.text
        response = await call("POST", paths[3] + "/revoke", "agent", {"expected_version": 1})
        assert response.status_code == 200, response.text
        with pytest.raises(ValueError, match="数量超过"):
            await rebuild(app.state.sessions, target_name, tenant_id=scenario["tenant"].bytes,
                          maximum_documents=1)
        with pytest.raises(ValueError, match="当前业务集合|前缀"):
            await rebuild(app.state.sessions, source.collection)
        async with asyncio.timeout(120):
            report = await rebuild(
                app.state.sessions, target_name, tenant_id=scenario["tenant"].bytes,
                maximum_documents=4, maximum_chunks=10, timeout_seconds=90,
            )
        assert report["documents_scanned"] == 4
        assert report["documents_restored"] == report["chunks_restored"] == 2
        process = await asyncio.create_subprocess_exec(
            sys.executable, "-m", "zhifu.modules.knowledge.recovery",
            "--target-collection", cli_name, "--tenant-id", str(scenario["tenant"]),
            "--maximum-documents", "4", "--maximum-chunks", "10",
            "--timeout-seconds", "30", "--output", ".work/knowledge-recovery-cli.json",
            stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE,
        )
        try:
            async with asyncio.timeout(45):
                output, error = await process.communicate()
        finally:
            if process.returncode is None:
                process.kill()
                await process.wait()
        assert process.returncode == 0, error.decode("utf-8", errors="replace")
        cli_report = json.loads(output)
        assert cli_report["status"] == "VERIFIED" and cli_report["documents_restored"] == 2
        assert cli_report["source_sha256"] == report["source_sha256"]
        with pytest.raises(ValueError, match="已存在"):
            await rebuild(app.state.sessions, target_name, tenant_id=scenario["tenant"].bytes)
        async with asyncio.timeout(90):
            concurrent_results = await asyncio.gather(*(
                rebuild(app.state.sessions, concurrent_name,
                        tenant_id=scenario["tenant"].bytes,
                        maximum_documents=4, maximum_chunks=10, timeout_seconds=60)
                for _ in range(2)
            ), return_exceptions=True)
        assert sum(isinstance(result, dict) for result in concurrent_results) == 1
        rejected = [result for result in concurrent_results if isinstance(result, ValueError)]
        assert len(rejected) == 1
        assert "已存在" in str(rejected[0]) or "其他进程创建" in str(rejected[0])
        async with retrieval_client(target) as recovered:
            app.state.knowledge = recovered
            path = f"/openapi/v1/stores/{scenario['store']}/knowledge/search"
            response = await call("POST", path, payload={"query": "键盘清洁断电"})
            assert response.status_code == 200, response.text
            assert {item["document_id"] for item in response.json()["items"]} == {
                value.rsplit("/", 1)[1] for value in paths[:2]
            }
            assert (await call("POST", path, "foreign", {"query": "键盘"})).status_code == 404
            _, source_digest = await snapshot(app.state.sessions, scenario["tenant"].bytes, 4)
            response = await call("POST", paths[0] + "/revoke", "agent", {"expected_version": 1})
            assert response.status_code == 200, response.text
            with pytest.raises(ValueError, match="源资料发生变化"):
                await verify_source(app.state.sessions, scenario["tenant"].bytes, 4, source_digest)
            response = await call("PUT", paths[1], "agent", {
                "expected_version": 1, "title": "修订资料", "text": "修订后等待重新发布。",
            })
            assert response.status_code == 200, response.text
            response = await call("POST", path, payload={"query": "键盘"})
            assert response.status_code == 200 and response.json()["items"] == []
        report["checks"] = {
            "draft_and_revoked_excluded": True, "foreign_tenant_denied": True,
            "current_revision_and_revocation_checked": True, "existing_target_denied": True,
            "active_target_denied": True, "document_limit_enforced": True,
            "cli_rebuild_verified": True,
            "same_target_concurrency_denied": True, "source_change_detected": True,
        }
        await asyncio.to_thread(
            Path(".work/knowledge-recovery-acceptance.json").write_text,
            json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8",
        )
    finally:
        app.state.knowledge = original
        # 仅删除本次随机生成的独立恢复集合，当前业务集合保持原样。
        def remove_target():
            for name in (target_name, cli_name, concurrent_name):
                with connection(source.model_copy(update={"collection": name})) as (
                    client, collection
                ):
                    if client.has_collection(collection):
                        client.drop_collection(collection)

        await asyncio.to_thread(remove_target)
        await cleanup_vectors(scenario["tenant"])