Add transactional file storage workflows, typed API contracts, recursive category handling, frontend component separation, and Files API coverage.
93 lines
2.8 KiB
Python
93 lines
2.8 KiB
Python
"""Tests for file workflow transaction compensation and path safety."""
|
|
|
|
from io import BytesIO
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
from agenteval.services.files import FileManagementService
|
|
from agenteval.storage.file_repository import FileRecordRepository
|
|
from agenteval.storage.file_storage import FileStorage, InvalidStorageNameError
|
|
from sqlmodel import Session, SQLModel, create_engine
|
|
from starlette.datastructures import Headers, UploadFile
|
|
|
|
|
|
@pytest.fixture()
|
|
def service_context(tmp_path: Path):
|
|
from agenteval.storage.db import ( # noqa: F401
|
|
EvalResultDB,
|
|
EvalRunDB,
|
|
EvalTargetDB,
|
|
FileCategoryDB,
|
|
FileRecordDB,
|
|
ScenarioDB,
|
|
TurnDB,
|
|
)
|
|
|
|
engine = create_engine(
|
|
f"sqlite:///{tmp_path / 'file_service.db'}",
|
|
connect_args={"check_same_thread": False},
|
|
)
|
|
SQLModel.metadata.create_all(engine)
|
|
session = Session(engine)
|
|
storage = FileStorage(tmp_path / "files")
|
|
service = FileManagementService(session, storage, {"txt"}, 1)
|
|
yield service, session, storage
|
|
session.close()
|
|
engine.dispose()
|
|
|
|
|
|
def _upload(content: bytes = b"content") -> UploadFile:
|
|
return UploadFile(
|
|
filename="test.txt",
|
|
file=BytesIO(content),
|
|
headers=Headers({"content-type": "text/plain"}),
|
|
)
|
|
|
|
|
|
async def test_upload_commit_failure_removes_physical_file(service_context, monkeypatch):
|
|
service, session, storage = service_context
|
|
|
|
def fail_commit():
|
|
raise RuntimeError("commit failed")
|
|
|
|
monkeypatch.setattr(session, "commit", fail_commit)
|
|
|
|
with pytest.raises(RuntimeError, match="commit failed"):
|
|
await service.upload(_upload())
|
|
|
|
assert list(storage.root.iterdir()) == []
|
|
assert FileRecordRepository(session).list_all() == []
|
|
|
|
|
|
def test_delete_commit_failure_restores_physical_file(service_context, monkeypatch):
|
|
service, session, storage = service_context
|
|
record = FileRecordRepository(session).create(
|
|
original_name="test.txt",
|
|
storage_name="stored.txt",
|
|
file_size=7,
|
|
mime_type="text/plain",
|
|
file_ext="txt",
|
|
)
|
|
session.commit()
|
|
physical_file = storage.path_for("stored.txt")
|
|
physical_file.write_bytes(b"content")
|
|
|
|
def fail_commit():
|
|
raise RuntimeError("commit failed")
|
|
|
|
monkeypatch.setattr(session, "commit", fail_commit)
|
|
|
|
with pytest.raises(RuntimeError, match="commit failed"):
|
|
service.delete_file(record.id)
|
|
|
|
assert physical_file.read_bytes() == b"content"
|
|
assert FileRecordRepository(session).get(record.id) is not None
|
|
|
|
|
|
@pytest.mark.parametrize("storage_name", ["../escape.txt", "folder/file.txt", ""])
|
|
def test_storage_rejects_names_outside_root(service_context, storage_name):
|
|
_, _, storage = service_context
|
|
|
|
with pytest.raises(InvalidStorageNameError):
|
|
storage.path_for(storage_name)
|