"""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}