deerflow-code/offline-backend-20260512/backend/tests/test_artifact_library.py

from __future__ import annotations

import os
from pathlib import Path
from types import SimpleNamespace

import pytest

from app.gateway import deps
from app.gateway.routers import artifacts


class _FakeThreadStore:
    def __init__(self, rows: list[dict]) -> None:
        self.rows = rows
        self.calls: list[dict] = []

    async def search(self, **kwargs) -> list[dict]:
        self.calls.append(kwargs)
        offset = int(kwargs.get("offset", 0))
        limit = int(kwargs.get("limit", 100))
        return self.rows[offset : offset + limit]


class _FakeAgentStore:
    async def list_by_ids(self, agent_ids: list[str], user_id: str) -> list[dict]:
        assert user_id == "owner"
        agents = {
            "analysis-agent": {
                "id": "analysis-agent",
                "name": "舆情分析智能体",
            }
        }
        return [agents[agent_id] for agent_id in agent_ids if agent_id in agents]


@pytest.mark.asyncio
async def test_artifact_library_aggregates_generated_non_media_files(tmp_path: Path, monkeypatch):
    thread_one = tmp_path / "thread-one" / "outputs"
    thread_two = tmp_path / "thread-two" / "outputs"
    thread_one.mkdir(parents=True)
    thread_two.mkdir(parents=True)

    report = thread_one / "研判报告.docx"
    table = thread_one / "数据汇总.xlsx"
    image = thread_one / "配图.png"
    audio = thread_two / "录音.mp3"
    notes = thread_two / "结论.md"
    archive = thread_two / "证据.zip"
    for path, content in (
        (report, b"doc"),
        (table, b"sheet"),
        (image, b"image"),
        (audio, b"audio"),
        (notes, b"# notes"),
        (archive, b"zip"),
    ):
        path.write_bytes(content)

    os.utime(report, (100, 100))
    os.utime(table, (200, 200))
    os.utime(notes, (300, 300))
    os.utime(archive, (400, 400))

    rows = [
        {
            "thread_id": "thread-one",
            "display_name": "战场环境综合研判",
            "metadata": {"agent_id": "analysis-agent"},
        },
        {
            "thread_id": "thread-two",
            "display_name": "生成发展报告",
            "metadata": {},
        },
    ]
    store = _FakeThreadStore(rows)

    class _Paths:
        def sandbox_outputs_dir(self, thread_id: str, *, user_id: str) -> Path:
            assert user_id == "owner"
            return tmp_path / thread_id / "outputs"

    async def _resolve_owner(_thread_id: str) -> str:
        return "owner"

    monkeypatch.setattr(artifacts, "get_paths", lambda: _Paths())
    monkeypatch.setattr(artifacts, "aresolve_path_user_id", _resolve_owner)

    items = await artifacts._scan_artifact_library(store)

    assert [item["name"] for item in items] == ["证据.zip", "结论.md", "数据汇总.xlsx", "研判报告.docx"]
    assert [item["file_kind"] for item in items] == ["archive", "text_data", "spreadsheet", "document"]
    assert [item["source_type"] for item in items] == ["general_qa", "general_qa", "agent", "agent"]
    assert items[-1]["thread_title"] == "战场环境综合研判"
    assert items[-1]["agent_id"] == "analysis-agent"
    assert all(item["name"] not in {"配图.png", "录音.mp3"} for item in items)
    assert store.calls[0]["exclude_system"] is True


@pytest.mark.asyncio
async def test_artifact_library_endpoint_filters_searches_and_pages(tmp_path: Path, monkeypatch):
    output_dir = tmp_path / "thread-one" / "outputs"
    agent_output_dir = tmp_path / "thread-agent" / "outputs"
    output_dir.mkdir(parents=True)
    agent_output_dir.mkdir(parents=True)
    (output_dir / "任务报告.md").write_text("report", encoding="utf-8")
    (output_dir / "任务数据.csv").write_text("a,b", encoding="utf-8")
    (agent_output_dir / "智能体结论.md").write_text("agent report", encoding="utf-8")

    store = _FakeThreadStore(
        [
            {
                "thread_id": "thread-one",
                "display_name": "任务研判对话",
                "metadata": {},
            },
            {
                "thread_id": "thread-agent",
                "display_name": "智能体研判对话",
                "metadata": {"agent_id": "analysis-agent"},
            },
        ]
    )

    class _Paths:
        def sandbox_outputs_dir(self, thread_id: str, *, user_id: str) -> Path:
            return tmp_path / thread_id / "outputs"

    async def _resolve_owner(_thread_id: str) -> str:
        return "owner"

    monkeypatch.setattr(artifacts, "get_paths", lambda: _Paths())
    monkeypatch.setattr(artifacts, "aresolve_path_user_id", _resolve_owner)
    monkeypatch.setattr(artifacts, "get_effective_user_id", lambda: "owner")
    monkeypatch.setattr(deps, "get_thread_store", lambda _request: store)
    monkeypatch.setattr(deps, "get_agent_store", lambda _request: _FakeAgentStore())
    artifacts._artifact_library_cache.clear()

    raw_endpoint = artifacts.list_artifact_library.__wrapped__
    response = await raw_endpoint(
        request=SimpleNamespace(),
        query="研判对话",
        file_kind="text_data",
        source_type="general_qa",
        offset=0,
        limit=1,
        refresh=True,
    )

    assert response.total == 1
    assert response.offset == 0
    assert response.limit == 1
    assert [item.name for item in response.items] == ["任务报告.md"]

    spreadsheet_response = await raw_endpoint(
        request=SimpleNamespace(),
        query="任务数据",
        file_kind="spreadsheet",
        source_type="all",
        offset=0,
        limit=50,
        refresh=False,
    )
    assert spreadsheet_response.total == 1
    assert spreadsheet_response.items[0].name == "任务数据.csv"

    agent_response = await raw_endpoint(
        request=SimpleNamespace(),
        query=None,
        file_kind="all",
        source_type="agent",
        offset=0,
        limit=50,
        refresh=False,
    )
    assert agent_response.total == 1
    assert agent_response.items[0].source_type == "agent"
    assert agent_response.items[0].agent_name == "舆情分析智能体"
    assert agent_response.items[0].name == "智能体结论.md"