xerrors--yuxi/backend/test/unit/knowledge/test_office_pdf_preview.py

from __future__ import annotations

from types import SimpleNamespace

import pytest

from yuxi.knowledge.base import KnowledgeBase
from yuxi.services.file_preview import MAX_BINARY_PREVIEW_SIZE_BYTES


class FakeKnowledgeBase(KnowledgeBase):
    @property
    def kb_type(self) -> str:
        return "fake"

    async def _create_kb_instance(self, kb_id: str, config: dict):
        return None

    async def _initialize_kb_instance(self, instance) -> None:
        pass

    async def index_file(self, kb_id: str, file_id: str, operator_id: str | None = None) -> dict:
        return {}

    async def update_content(self, kb_id: str, file_ids: list[str], params: dict | None = None) -> list[dict]:
        return []

    async def aquery(self, query_text: str, kb_id: str, **kwargs) -> list[dict]:
        return []

    def get_query_params_config(self, kb_id: str, **kwargs) -> dict:
        return {"options": []}

    async def delete_file(self, kb_id: str, file_id: str) -> None:
        pass

    async def get_file_basic_info(self, kb_id: str, file_id: str) -> dict:
        return {}

    async def get_file_content(self, kb_id: str, file_id: str) -> dict:
        return {}

    async def get_file_info(self, kb_id: str, file_id: str) -> dict:
        return {}

    async def _save_metadata(self) -> None:
        pass


class FakeMinioClient:
    KB_BUCKETS = {"parsed": "knowledgebases"}

    def __init__(self) -> None:
        self.objects: dict[tuple[str, str], bytes] = {}

    async def astat_file(self, bucket_name: str, object_name: str) -> int | None:
        content = self.objects.get((bucket_name, object_name))
        return len(content) if content is not None else None

    async def adownload_file(self, bucket_name: str, object_name: str) -> bytes:
        return self.objects[(bucket_name, object_name)]

    async def aupload_file(
        self,
        bucket_name: str,
        object_name: str,
        data: bytes,
        content_type: str | None = None,
    ) -> SimpleNamespace:
        assert content_type == "application/pdf"
        self.objects[(bucket_name, object_name)] = data
        return SimpleNamespace(url=f"http://localhost:9000/{bucket_name}/{object_name}")


def make_kb(tmp_path) -> FakeKnowledgeBase:
    kb = FakeKnowledgeBase(str(tmp_path))
    kb.databases_meta["db1"] = {"metadata": {}}
    kb.test_file_meta = {
        "file_id": "file1",
        "kb_id": "db1",
        "filename": "demo.docx",
        "path": "minio://knowledgebases/db1/upload/demo.docx",
        "markdown_file": "minio://knowledgebases/db1/parsed/file1.md",
        "status": "parsed",
    }

    async def load_file_meta(kb_id: str, file_id: str, *, refresh: bool = False) -> dict:
        assert kb_id == "db1"
        assert file_id == "file1"
        return kb.test_file_meta

    kb._load_file_meta = load_file_meta
    return kb


def test_office_file_entry_exposes_logical_file_availability(tmp_path) -> None:
    kb = make_kb(tmp_path)

    entry = kb._knowledge_file_entry("db1", "file1", kb.test_file_meta)

    assert entry["has_original_file"] is True
    assert entry["has_parsed_markdown"] is True
    assert "preview_modes" not in entry
    assert "default_preview_mode" not in entry


@pytest.mark.asyncio
async def test_read_office_pdf_preview_converts_and_caches_pdf(tmp_path, monkeypatch) -> None:
    kb = make_kb(tmp_path)
    minio_client = FakeMinioClient()
    minio_client.objects[("knowledgebases", "db1/upload/demo.docx")] = b"office"
    minio_client.objects[("knowledgebases", "db1/parsed/file1.md")] = b"# parsed"

    convert_calls = 0

    async def fake_convert(filename: str, content: bytes) -> bytes:
        nonlocal convert_calls
        convert_calls += 1
        assert filename == "demo.docx"
        assert content == b"office"
        return b"%PDF-1.4\nconverted"

    monkeypatch.setattr("yuxi.storage.minio.get_minio_client", lambda: minio_client)
    monkeypatch.setattr("yuxi.knowledge.base.convert_office_to_pdf", fake_convert)

    response = await kb.read_file_preview("db1", "file1")
    cached_response = await kb.read_file_preview("db1", "file1")

    assert response["preview_type"] == "pdf"
    assert response["supported"] is True
    assert response["binary"] is True
    assert response["content"] == b"%PDF-1.4\nconverted"
    assert response["media_type"] == "application/pdf"
    assert cached_response["content"] == b"%PDF-1.4\nconverted"
    assert minio_client.objects[("knowledgebases", "db1/preview/file1.pdf")] == b"%PDF-1.4\nconverted"
    assert convert_calls == 1


@pytest.mark.asyncio
async def test_non_docx_pptx_office_files_do_not_get_pdf_preview(tmp_path, monkeypatch) -> None:
    kb = make_kb(tmp_path)
    kb.test_file_meta["filename"] = "demo.xlsx"
    minio_client = FakeMinioClient()
    minio_client.objects[("knowledgebases", "db1/upload/demo.docx")] = b"PK\x03\x04excel"
    monkeypatch.setattr("yuxi.storage.minio.get_minio_client", lambda: minio_client)

    entry = kb._knowledge_file_entry("db1", "file1", kb.test_file_meta)
    response = await kb.read_file_preview("db1", "file1")

    assert entry["has_original_file"] is True
    assert entry["has_parsed_markdown"] is True
    assert response["preview_type"] == "unsupported"
    assert response["supported"] is False


@pytest.mark.asyncio
async def test_read_file_preview_rejects_large_original_before_download(tmp_path, monkeypatch) -> None:
    kb = make_kb(tmp_path)
    kb.test_file_meta["filename"] = "large.pdf"
    kb.test_file_meta["size"] = MAX_BINARY_PREVIEW_SIZE_BYTES + 1

    async def fail_read(_path: str) -> bytes:
        raise AssertionError("large preview should not download file content")

    monkeypatch.setattr(kb, "_read_minio_bytes", fail_read)

    response = await kb.read_file_preview("db1", "file1")

    assert response["preview_type"] == "unsupported"
    assert response["supported"] is False
    assert response["limit"] == MAX_BINARY_PREVIEW_SIZE_BYTES