PlantAI-rag-new/tooling/scripts/test_docgen_import.py

"""tooling/scripts/test_docgen_import.py —— 批量导入工具单测（P2，本地运行，不连 vstore）。

用 httpx.MockTransport 注入假 vstore（ingest / rebuild-index 两个端点），
覆盖：metadata.csv 映射、并发 ingest、--dry-run 不删目录、部分失败不删目录、
成功后删除源明文目录、--keep-source 保留。
"""
from __future__ import annotations

import sys
from pathlib import Path

import httpx
import pytest

_RealClient = httpx.Client  # 捕获原实现（必须在 monkeypatch httpx.Client 之前）

SCRIPT = Path(__file__).resolve().parent / "docgen_import.py"

# 把脚本加载为模块（tooling 不进 sdk 编译链，测试独立跑）
import importlib.util as _ilu

_spec = _ilu.spec_from_file_location("docgen_import", SCRIPT)
docgen_import = _ilu.module_from_spec(_spec)
_spec.loader.exec_module(docgen_import)


# ---------------------------------------------------------------------------
# 假 vstore：记录 ingest 请求，成功返回 document_id
# ---------------------------------------------------------------------------

class FakeVstore:
    def __init__(self):
        self.ingested: list[dict] = []
        self.rebuilds = 0

    def handler(self, request: httpx.Request) -> httpx.Response:
        if request.url.path == "/rebuild-index":
            self.rebuilds += 1
            return httpx.Response(200, json={"rebuilt": True})
        if request.url.path == "/ingest":
            # multipart body 里取 file 名（解析 bytes 中 filename=）
            body = request.content
            name = "unknown"
            if b'filename="' in body:
                name = body.split(b'filename="')[1].split(b'"')[0].decode("utf-8", "replace")
            self.ingested.append(
                {"name": name, "fields": dict(request.url.params) if False else {}}
            )
            # multipart 表单字段在 body 里；粗略提取 as_template 便于断言
            as_template = b"as_template" in body
            self.ingested[-1]["as_template_field_present"] = as_template
            return httpx.Response(200, json={"document_id": f"doc-{len(self.ingested)}"})
        return httpx.Response(404, json={"detail": "not found"})


def _client(fake: FakeVstore) -> httpx.Client:
    return _RealClient(transport=httpx.MockTransport(fake.handler), base_url="http://vstore.test")


def _make_source_dir(tmp_path: Path, files: dict[str, str] | None = None,
                     csv_content: str | None = None) -> Path:
    src = tmp_path / "src"
    src.mkdir()
    for rel, content in (files or {"a.md": "内容A"}).items():
        p = src / rel
        p.parent.mkdir(parents=True, exist_ok=True)
        p.write_text(content, encoding="utf-8")
    if csv_content is not None:
        (src / "metadata.csv").write_text(csv_content, encoding="utf-8")
    return src


def _run_main(argv: list[str], monkeypatch) -> int:
    import docgen_import as di

    monkeypatch.setattr(sys, "argv", ["docgen_import.py", *argv])
    try:
        return di.main()
    except SystemExit as exc:
        code = exc.code
        if code is None:
            return 0
        if isinstance(code, int):
            return code
        print(code, file=sys.stderr)  # fail() 打印的信息
        return 2


def test_dry_run_lists_files_and_keeps_source(tmp_path, monkeypatch):
    src = _make_source_dir(tmp_path, {"a.md": "内容A", "sub/b.txt": "内容B"},
                           csv_content="file,name,as_template\na.md,资料甲,false\n")
    code = _run_main(
        ["--dir", str(src), "--dry-run", "--vstore", "http://vstore.test"], monkeypatch
    )
    assert code == 0
    assert src.is_dir()  # dry-run 不删目录
    assert (src / "a.md").exists()


def test_ingest_ok_and_source_removed(tmp_path, monkeypatch):
    src = _make_source_dir(tmp_path, {"a.md": "内容A"})
    fake = FakeVstore()
    monkeypatch.setattr(docgen_import.httpx, "Client", lambda **kw: _client(fake))
    code = _run_main(["--dir", str(src), "--vstore", "http://vstore.test",
                      "--workers", "2"], monkeypatch)
    assert code == 0
    assert len(fake.ingested) == 1
    assert fake.ingested[0]["name"] == "a.md"
    assert fake.rebuilds == 1
    assert not src.exists()  # 导入完成删除源明文目录


def test_keep_source_retains_dir(tmp_path, monkeypatch):
    src = _make_source_dir(tmp_path, {"a.md": "内容A"})
    fake = FakeVstore()
    monkeypatch.setattr(docgen_import.httpx, "Client", lambda **kw: _client(fake))
    code = _run_main(["--dir", str(src), "--keep-source"], monkeypatch)
    assert code == 0
    assert src.is_dir()
    assert (src / "a.md").exists()


def test_metadata_csv_mapping(tmp_path, monkeypatch):
    src = _make_source_dir(tmp_path, {"a.md": "内容A"},
                           csv_content="file,name,as_template\na.md,资料甲,false\n")
    fake = FakeVstore()
    monkeypatch.setattr(docgen_import.httpx, "Client", lambda **kw: _client(fake))
    code = _run_main(["--dir", str(src)], monkeypatch)
    assert code == 0
    # as_template=false 会以 "false" 字符串出现在 multipart 字段里
    body_fields = fake.ingested[0]
    assert body_fields["name"] == "a.md"


def test_partial_failure_keeps_source(tmp_path, monkeypatch):
    src = _make_source_dir(tmp_path, {"a.md": "内容A", "b.md": "内容B"})
    fake = FakeVstore()

    def flaky(request: httpx.Request) -> httpx.Response:
        if b'filename="b.md"' in request.content:
            return httpx.Response(500, json={"detail": "boom"})
        return fake.handler(request)

    monkeypatch.setattr(
        docgen_import.httpx, "Client",
        lambda **kw: _RealClient(transport=httpx.MockTransport(flaky), base_url="http://vstore.test"),
    )
    code = _run_main(["--dir", str(src)], monkeypatch)
    assert code == 2  # 有失败 → 非零退出
    assert src.is_dir()  # 失败不删源目录
    assert fake.rebuilds == 0  # 失败不重建索引


def test_missing_dir_fails(tmp_path, monkeypatch):
    code = _run_main(["--dir", str(tmp_path / "nope")], monkeypatch)
    assert code == 2


def test_collect_skips_unsupported_and_metadata_csv(tmp_path):
    src = _make_source_dir(tmp_path, {"a.md": "x", "b.exe": "x"},
                           csv_content="file,name\na.md,甲\n")
    files = docgen_import.collect_files(src)
    assert [p.name for p in files] == ["a.md"]