vent_agent_2.0/tests/test_chat_uploads.py

"""附件上传：多份资料一次传入、数量上限、与 1.5 单份字段的合并。

背景：风量研判至少要「配风计划 + 测风报表」两份文件，而 1.5 契约里只有单个 `file`
字段——所以加了 `files` 多附件字段，且两者必须合并处理。
"""
from __future__ import annotations

from io import BytesIO
from typing import Any

import pytest
from fastapi import HTTPException
from starlette.datastructures import UploadFile

from api import chat_routes as cr
from agents import workspace


@pytest.fixture
def isolated_sandbox(tmp_path, monkeypatch):
    """把会话沙箱挪到临时目录，避免测试往仓库 data/workspaces 里写文件。"""
    monkeypatch.setattr(workspace, "session_root", lambda user, session: tmp_path / user / session)
    return tmp_path


def _upload(name: str, payload: bytes = b"data") -> UploadFile:
    return UploadFile(filename=name, file=BytesIO(payload))


# ---- 多附件 ---------------------------------------------------------------

@pytest.mark.asyncio
async def test_save_uploads_keeps_all_attachments(isolated_sandbox):
    saved = await cr._save_uploads(
        [_upload("配风计划.pdf"), _upload("测风报表.xlsx"), _upload("监控.csv")], "u", "s")
    assert [name for name, _ in saved] == ["配风计划.pdf", "测风报表.xlsx", "监控.csv"]
    paths = [path for _, path in saved]
    assert all(path.startswith("/workspace/uploads/") for path in paths)
    assert len(set(paths)) == 3, "同名文件不能互相覆盖"
    for path in paths:
        real = isolated_sandbox / "u" / "s" / path.lstrip("/")
        assert real.is_file() and real.read_bytes() == b"data"


@pytest.mark.asyncio
async def test_save_uploads_rejects_too_many_files(isolated_sandbox, monkeypatch):
    monkeypatch.setattr(cr, "MAX_UPLOAD_FILES", 2)
    with pytest.raises(HTTPException) as excinfo:
        await cr._save_uploads([_upload(f"f{i}.txt") for i in range(3)], "u", "s")
    assert excinfo.value.status_code == 400
    assert excinfo.value.detail["code"] == "too_many_files"
    # 先卡数量再落盘：被拒时不应留下半截文件
    assert not list((isolated_sandbox / "u" / "s" / "workspace" / "uploads").glob("*"))


@pytest.mark.asyncio
async def test_save_uploads_accepts_empty_list(isolated_sandbox):
    assert await cr._save_uploads([], "u", "s") == []


def test_inject_all_lists_every_upload():
    ref: dict[str, Any] = {"reference": None, "forced_skill": None}
    single = cr._inject_all("看一下", "u", "full", ref, [("a.xlsx", "/workspace/uploads/a.xlsx")])
    assert "「a.xlsx」" in single and "1 个文件" not in single

    multi = cr._inject_all("看一下", "u", "full", ref, [
        ("a.pdf", "/workspace/uploads/a.pdf"), ("b.xlsx", "/workspace/uploads/b.xlsx")])
    assert "2 个文件" in multi
    for token in ("a.pdf", "b.xlsx", "/workspace/uploads/a.pdf", "/workspace/uploads/b.xlsx"):
        assert token in multi
    assert "不要在文件系统里搜索别处的资料" in multi


@pytest.mark.asyncio
async def test_chat_merges_single_file_field_and_files_field(monkeypatch):
    """1.5 的 file 与 2.0 的 files 必须合并处理，不能只认一个。"""
    captured: dict[str, Any] = {}

    async def fake_pipeline(**kwargs):
        captured.update(kwargs)
        return "stub"

    monkeypatch.setattr(cr, "_run_chat_pipeline", fake_pipeline)
    result = await cr.chat(message="研判", session_id=None, thread_id=None, mode=None,
                           file=_upload("单份.pdf"),
                           files=[_upload("甲.xlsx"), _upload("乙.xlsx")],
                           user_name=None, request=None, request_id=None)
    assert result == "stub"
    names = [item.filename for item in captured["attachments"]]
    assert names == ["单份.pdf", "甲.xlsx", "乙.xlsx"]
    assert captured["user_name"] == "default"


@pytest.mark.asyncio
async def test_chat_without_attachments_passes_empty_list(monkeypatch):
    captured: dict[str, Any] = {}

    async def fake_pipeline(**kwargs):
        captured.update(kwargs)
        return "stub"

    monkeypatch.setattr(cr, "_run_chat_pipeline", fake_pipeline)
    await cr.chat(message="随便聊聊", session_id=None, thread_id=None, mode=None,
                  file=None, files=[], user_name=None, request=None, request_id=None)
    assert captured["attachments"] == []