ai-video-fullstack/backend/routes/test_assets.py

"""Persistent image fixtures for repeatable assistant tests."""

from __future__ import annotations

import asyncio
from pathlib import Path
from uuid import uuid4

from db.models import TestAsset, TestCase, TestRunCase
from db.session import get_session
from fastapi import APIRouter, Depends, File, HTTPException, UploadFile
from fastapi.responses import Response
from loguru import logger
from pydantic import BaseModel
from services.auth import require_admin
from services.test_assets import delete_test_image, read_test_image, store_test_image
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

import settings


router = APIRouter(
    prefix="/api/test-assets",
    tags=["test-assets"],
    dependencies=[Depends(require_admin)],
)


class TestImageAssetOut(BaseModel):
    id: str
    filename: str
    width: int
    height: int
    sizeBytes: int
    contentUrl: str


def _new_id() -> str:
    return f"ta_{uuid4().hex}"


def _safe_filename(value: str | None) -> str:
    name = Path(value or "image.jpg").name.strip()
    return (name or "image.jpg")[:255]


def _out(row: TestAsset) -> TestImageAssetOut:
    return TestImageAssetOut(
        id=row.id,
        filename=row.original_name,
        width=row.width,
        height=row.height,
        sizeBytes=row.size_bytes,
        contentUrl=f"/api/test-assets/{row.id}/content",
    )


def _contains_asset_id(value: object, asset_id: str) -> bool:
    if isinstance(value, dict):
        if value.get("assetId") == asset_id or value.get("asset_id") == asset_id:
            return True
        return any(_contains_asset_id(item, asset_id) for item in value.values())
    if isinstance(value, list):
        return any(_contains_asset_id(item, asset_id) for item in value)
    return False


async def _is_referenced(session: AsyncSession, asset_id: str) -> bool:
    definitions = (
        await session.execute(select(TestCase.definition))
    ).scalars().all()
    if any(_contains_asset_id(value, asset_id) for value in definitions):
        return True
    snapshots = (
        await session.execute(select(TestRunCase.case_snapshot))
    ).scalars().all()
    return any(_contains_asset_id(value, asset_id) for value in snapshots)


@router.post("/images", response_model=TestImageAssetOut)
async def upload_test_image(
    file: UploadFile = File(...),
    session: AsyncSession = Depends(get_session),
):
    data = await file.read(settings.INPUT_IMAGE_MAX_BYTES + 1)
    if len(data) > settings.INPUT_IMAGE_MAX_BYTES:
        limit_mb = settings.INPUT_IMAGE_MAX_BYTES // 1024 // 1024
        raise HTTPException(413, f"图片不能超过 {limit_mb} MB")
    if not data:
        raise HTTPException(400, "图片文件不能为空")

    asset_id = _new_id()
    try:
        stored = await asyncio.to_thread(store_test_image, asset_id, data)
    except ValueError as exc:
        raise HTTPException(400, str(exc)) from exc
    except Exception as exc:
        raise HTTPException(503, "图片暂时无法上传，请确认对象存储可用") from exc

    row = TestAsset(
        id=asset_id,
        original_name=_safe_filename(file.filename),
        storage_key=stored.storage_key,
        mime_type=stored.mime_type,
        width=stored.width,
        height=stored.height,
        size_bytes=stored.size_bytes,
    )
    session.add(row)
    try:
        await session.commit()
    except Exception:
        await session.rollback()
        try:
            await asyncio.to_thread(delete_test_image, stored.storage_key)
        except Exception:
            logger.warning(f"回滚测试图片上传时清理对象失败: {stored.storage_key}")
        raise
    await session.refresh(row)
    return _out(row)


@router.get("/{asset_id}", response_model=TestImageAssetOut)
async def get_test_asset(
    asset_id: str,
    session: AsyncSession = Depends(get_session),
):
    row = await session.get(TestAsset, asset_id)
    if row is None:
        raise HTTPException(404, "测试图片不存在")
    return _out(row)


@router.get("/{asset_id}/content")
async def get_test_asset_content(
    asset_id: str,
    session: AsyncSession = Depends(get_session),
):
    row = await session.get(TestAsset, asset_id)
    if row is None:
        raise HTTPException(404, "测试图片不存在")
    try:
        data = await asyncio.to_thread(read_test_image, row.storage_key)
    except Exception as exc:
        raise HTTPException(503, "测试图片暂时无法读取") from exc
    return Response(
        content=data,
        media_type=row.mime_type,
        headers={"Cache-Control": "private, max-age=3600"},
    )


@router.delete("/{asset_id}")
async def delete_test_asset(
    asset_id: str,
    session: AsyncSession = Depends(get_session),
):
    row = await session.get(TestAsset, asset_id)
    if row is None:
        raise HTTPException(404, "测试图片不存在")
    if await _is_referenced(session, asset_id):
        raise HTTPException(409, "测试图片仍被用例或运行记录引用")

    storage_key = row.storage_key
    await session.delete(row)
    await session.commit()
    try:
        await asyncio.to_thread(delete_test_image, storage_key)
    except Exception:
        logger.warning(f"测试图片数据已删除，但对象存储清理失败: {storage_key}")
    return {"ok": True}