AgentEvalTool/tests/integration/test_files_api.py

"""Integration tests for the file management API."""

from pathlib import Path

import pytest
from agenteval.config import get_settings
from agenteval.web.app import app
from agenteval.web.deps import get_db
from fastapi.testclient import TestClient
from sqlmodel import Session, SQLModel, create_engine


@pytest.fixture()
def files_client(tmp_path: Path, monkeypatch):
    from agenteval.storage.db import (  # noqa: F401
        EvalResultDB,
        EvalRunDB,
        EvalTargetDB,
        FileCategoryDB,
        FileRecordDB,
        ScenarioDB,
        TurnDB,
    )
    from agenteval.web.routers import files as files_module

    engine = create_engine(
        f"sqlite:///{tmp_path / 'files_api.db'}",
        connect_args={"check_same_thread": False},
    )
    SQLModel.metadata.create_all(engine)
    session = Session(engine)
    files_dir = tmp_path / "files"
    files_dir.mkdir()

    monkeypatch.setattr(files_module, "FILES_DIR", files_dir)
    settings = get_settings()
    monkeypatch.setattr(settings, "allowed_extensions", "txt,json")
    monkeypatch.setattr(settings, "max_upload_size_mb", 1)

    def override_get_db():
        yield session

    app.dependency_overrides[get_db] = override_get_db
    with TestClient(app) as client:
        yield client, session, files_dir

    app.dependency_overrides.clear()
    session.close()
    engine.dispose()


def _create_category(client: TestClient, name: str, parent_id: str | None = None) -> dict:
    response = client.post(
        "/api/files/categories",
        json={"name": name, "parent_id": parent_id},
    )
    assert response.status_code == 200
    return response.json()


def test_upload_config_uses_server_settings(files_client):
    client, _, _ = files_client

    response = client.get("/api/files/config")

    assert response.status_code == 200
    assert response.json() == {
        "allowed_extensions": ["json", "txt"],
        "max_upload_size_mb": 1,
    }


def test_category_crud_returns_domain_fields_and_deep_tree(files_client):
    client, _, _ = files_client
    root = _create_category(client, "根分类")
    child = _create_category(client, "二级", root["id"])
    leaf = _create_category(client, "三级", child["id"])

    tree = client.get("/api/files/categories").json()

    assert tree[0]["id"] == root["id"]
    assert tree[0]["name"] == "根分类"
    assert tree[0]["children"][0]["id"] == child["id"]
    assert tree[0]["children"][0]["children"][0]["id"] == leaf["id"]

    updated = client.put(
        f"/api/files/categories/{leaf['id']}",
        json={"name": "三级分类"},
    )
    assert updated.status_code == 200
    assert updated.json()["name"] == "三级分类"


def test_create_category_rejects_missing_parent(files_client):
    client, _, _ = files_client

    response = client.post(
        "/api/files/categories",
        json={"name": "孤立分类", "parent_id": "missing"},
    )

    assert response.status_code == 400
    assert response.json()["detail"] == "父分类不存在"


def test_list_files_includes_all_descendant_categories(files_client):
    client, _, _ = files_client
    root = _create_category(client, "根分类")
    child = _create_category(client, "二级", root["id"])
    leaf = _create_category(client, "三级", child["id"])

    upload = client.post(
        "/api/files/upload",
        data={"category_id": leaf["id"]},
        files={"file": ("deep.txt", b"deep content", "text/plain")},
    )
    assert upload.status_code == 200

    response = client.get("/api/files", params={"category_id": root["id"]})

    assert response.status_code == 200
    assert [item["original_name"] for item in response.json()] == ["deep.txt"]


def test_upload_download_and_delete_file(files_client):
    client, _, files_dir = files_client

    uploaded = client.post(
        "/api/files/upload",
        files={"file": ("notes.txt", b"hello", "text/plain")},
    )

    assert uploaded.status_code == 200
    record = uploaded.json()
    assert record["original_name"] == "notes.txt"
    assert record["file_size"] == 5
    assert record["storage_directory"] == "data/files"
    assert client.get("/api/files").json()[0]["storage_directory"] == "data/files"
    assert len(list(files_dir.iterdir())) == 1

    downloaded = client.get(f"/api/files/{record['id']}/download")
    assert downloaded.status_code == 200
    assert downloaded.content == b"hello"

    deleted = client.delete(f"/api/files/{record['id']}")
    assert deleted.status_code == 200
    assert list(files_dir.iterdir()) == []
    assert client.get("/api/files").json() == []


def test_upload_rejects_invalid_category_and_extension(files_client):
    client, _, files_dir = files_client

    invalid_category = client.post(
        "/api/files/upload",
        data={"category_id": "missing"},
        files={"file": ("notes.txt", b"hello", "text/plain")},
    )
    invalid_extension = client.post(
        "/api/files/upload",
        files={"file": ("notes.exe", b"hello", "application/octet-stream")},
    )

    assert invalid_category.status_code == 400
    assert invalid_category.json()["detail"] == "分类不存在"
    assert invalid_extension.status_code == 400
    assert list(files_dir.iterdir()) == []


def test_upload_rejects_oversized_file_without_disk_artifact(files_client):
    client, _, files_dir = files_client

    response = client.post(
        "/api/files/upload",
        files={"file": ("large.txt", b"x" * (1024 * 1024 + 1), "text/plain")},
    )

    assert response.status_code == 413
    assert list(files_dir.iterdir()) == []
    assert client.get("/api/files").json() == []


def test_delete_category_removes_descendant_files(files_client):
    client, _, files_dir = files_client
    root = _create_category(client, "根分类")
    child = _create_category(client, "子分类", root["id"])
    uploaded = client.post(
        "/api/files/upload",
        data={"category_id": child["id"]},
        files={"file": ("child.txt", b"content", "text/plain")},
    )
    assert uploaded.status_code == 200
    assert len(list(files_dir.iterdir())) == 1

    response = client.delete(f"/api/files/categories/{root['id']}")

    assert response.status_code == 200
    assert client.get("/api/files/categories").json() == []
    assert client.get("/api/files").json() == []
    assert list(files_dir.iterdir()) == []