AgentEvalTool/tests/unit/test_file_repository.py

"""Unit tests for FileCategoryRepository and FileRecordRepository."""

import uuid
from pathlib import Path

import pytest
from sqlmodel import Session, SQLModel, create_engine

from agenteval.storage.file_repository import FileCategoryRepository, FileRecordRepository


@pytest.fixture()
def file_db_session(tmp_path: Path):
    """Isolated SQLite session with file management tables."""
    from agenteval.storage.db import (  # noqa: F401
        EvalResultDB, EvalRunDB, EvalTargetDB, FileCategoryDB, FileRecordDB, ScenarioDB, TurnDB,
    )
    engine = create_engine(
        f"sqlite:///{tmp_path / 'file_test.db'}",
        connect_args={"check_same_thread": False},
    )
    SQLModel.metadata.create_all(engine)
    session = Session(engine)
    try:
        yield session
    finally:
        session.close()
        engine.dispose()


@pytest.fixture()
def fake_files_dir(tmp_path: Path, monkeypatch):
    """Redirect FILES_DIR to a temp directory so tests don't touch data/files."""
    import agenteval.storage.db as db_mod
    import agenteval.storage.file_repository as repo_mod

    fake_dir = tmp_path / "files"
    fake_dir.mkdir()
    monkeypatch.setattr(db_mod, "FILES_DIR", fake_dir)
    monkeypatch.setattr(repo_mod, "FILES_DIR", fake_dir)
    return fake_dir


# ── FileCategoryRepository ────────────────────────────────────────────────

def test_create_root_category(file_db_session):
    repo = FileCategoryRepository(file_db_session)
    cat = repo.create("文件一类")
    assert cat.id is not None
    assert cat.name == "文件一类"
    assert cat.parent_id is None


def test_create_child_category(file_db_session):
    repo = FileCategoryRepository(file_db_session)
    root = repo.create("根分类")
    child = repo.create("子分类", parent_id=root.id)
    assert child.parent_id == root.id


def test_list_all_categories(file_db_session):
    repo = FileCategoryRepository(file_db_session)
    repo.create("A")
    repo.create("B")
    all_cats = repo.list_all()
    assert len(all_cats) == 2


def test_get_tree_structure(file_db_session):
    repo = FileCategoryRepository(file_db_session)
    root = repo.create("root")
    repo.create("child1", parent_id=root.id)
    repo.create("child2", parent_id=root.id)
    tree = repo.get_tree()
    assert len(tree) == 1
    assert tree[0]["key"] == root.id
    assert len(tree[0]["children"]) == 2


def test_get_tree_flat(file_db_session):
    repo = FileCategoryRepository(file_db_session)
    repo.create("A")
    repo.create("B")
    tree = repo.get_tree()
    assert len(tree) == 2
    assert all(len(node["children"]) == 0 for node in tree)


def test_update_category_name(file_db_session):
    repo = FileCategoryRepository(file_db_session)
    cat = repo.create("旧名称")
    updated = repo.update(cat.id, "新名称")
    assert updated is not None
    assert updated.name == "新名称"


def test_update_nonexistent_returns_none(file_db_session):
    repo = FileCategoryRepository(file_db_session)
    assert repo.update("nonexistent-id", "name") is None


def test_get_nonexistent_returns_none(file_db_session):
    repo = FileCategoryRepository(file_db_session)
    assert repo.get("no-such-id") is None


def test_delete_category(file_db_session):
    repo = FileCategoryRepository(file_db_session)
    cat = repo.create("删除目标")
    assert repo.delete(cat.id) is True
    assert repo.get(cat.id) is None


def test_delete_nonexistent_returns_false(file_db_session):
    repo = FileCategoryRepository(file_db_session)
    assert repo.delete("ghost-id") is False


def test_delete_cascades_to_children(file_db_session, fake_files_dir):
    cat_repo = FileCategoryRepository(file_db_session)
    file_repo = FileRecordRepository(file_db_session)

    root = cat_repo.create("父")
    child = cat_repo.create("子", parent_id=root.id)

    # Create a file in child
    file_repo.create("f.txt", "f-storage.txt", 10, "text/plain", "txt", category_id=child.id)

    assert cat_repo.delete(root.id) is True
    assert cat_repo.get(root.id) is None
    assert cat_repo.get(child.id) is None
    assert len(file_repo.list_all()) == 0


# ── FileRecordRepository ──────────────────────────────────────────────────

def test_create_file_record(file_db_session):
    repo = FileRecordRepository(file_db_session)
    rec = repo.create("test.txt", "stored-uuid.txt", 128, "text/plain", "txt")
    assert rec.id is not None
    assert rec.original_name == "test.txt"
    assert rec.file_size == 128


def test_create_file_with_category(file_db_session):
    cat_repo = FileCategoryRepository(file_db_session)
    file_repo = FileRecordRepository(file_db_session)
    cat = cat_repo.create("分类")
    rec = file_repo.create("data.json", "stored.json", 256, "application/json", "json", category_id=cat.id)
    assert rec.category_id == cat.id


def test_list_files_all(file_db_session):
    repo = FileRecordRepository(file_db_session)
    repo.create("a.txt", "a-s.txt", 1, "text/plain", "txt")
    repo.create("b.txt", "b-s.txt", 2, "text/plain", "txt")
    assert len(repo.list_all()) == 2


def test_list_files_by_category(file_db_session):
    cat_repo = FileCategoryRepository(file_db_session)
    file_repo = FileRecordRepository(file_db_session)
    cat_a = cat_repo.create("A")
    cat_b = cat_repo.create("B")
    file_repo.create("in_a.txt", "s1.txt", 1, "text/plain", "txt", category_id=cat_a.id)
    file_repo.create("in_b.txt", "s2.txt", 2, "text/plain", "txt", category_id=cat_b.id)
    file_repo.create("no_cat.txt", "s3.txt", 3, "text/plain", "txt")

    only_a = file_repo.list_all(category_id=cat_a.id)
    assert len(only_a) == 1
    assert only_a[0].original_name == "in_a.txt"


def test_list_files_includes_subcategory_files(file_db_session):
    cat_repo = FileCategoryRepository(file_db_session)
    file_repo = FileRecordRepository(file_db_session)
    root = cat_repo.create("root")
    child = cat_repo.create("child", parent_id=root.id)
    file_repo.create("in_root.txt", "sr.txt", 1, "text/plain", "txt", category_id=root.id)
    file_repo.create("in_child.txt", "sc.txt", 2, "text/plain", "txt", category_id=child.id)

    files = file_repo.list_all(category_id=root.id)
    assert len(files) == 2


def test_get_file_record(file_db_session):
    repo = FileRecordRepository(file_db_session)
    rec = repo.create("get_me.txt", "gm.txt", 10, "text/plain", "txt")
    fetched = repo.get(rec.id)
    assert fetched is not None
    assert fetched.id == rec.id


def test_get_nonexistent_file_returns_none(file_db_session):
    repo = FileRecordRepository(file_db_session)
    assert repo.get("no-such-id") is None


def test_delete_file_record(file_db_session):
    repo = FileRecordRepository(file_db_session)
    rec = repo.create("del.txt", "del-s.txt", 5, "text/plain", "txt")
    assert repo.delete(rec.id) is True
    assert repo.get(rec.id) is None


def test_delete_removes_physical_file(file_db_session, fake_files_dir):
    repo = FileRecordRepository(file_db_session)
    storage_name = f"{uuid.uuid4()}.txt"
    physical = fake_files_dir / storage_name
    physical.write_text("content")
    assert physical.exists()

    rec = repo.create("orig.txt", storage_name, 7, "text/plain", "txt")
    repo.delete(rec.id)
    assert not physical.exists()


def test_delete_nonexistent_file_returns_false(file_db_session):
    repo = FileRecordRepository(file_db_session)
    assert repo.delete("ghost-id") is False