geMoldInsight/tests/test_batch_status_pg.py

"""批次 2（D7）回归测试：批量任务聚合查询以 PG 为单一事实源。

覆盖：
- GET /api/batch/{batch_id} 按 ProcessingTask.batch_id 聚合（此前依赖 Redis batch key + 进程内存降级）
- 归属校验：他人批次 403、不存在 404、所有者 200 + 聚合数字正确
"""
import pytest
from fastapi import FastAPI
from httpx import AsyncClient, ASGITransport
from sqlalchemy.ext.asyncio import async_sessionmaker, AsyncSession

from moldinsight.api.batch_router import router as batch_router
from shared.database.database import get_db_session
from shared.models.identity import User
from moldinsight.models import STPFile, ProcessingTask
from shared.services.auth_service import get_current_active_user


@pytest.fixture(scope="function")
async def batch_client(async_engine, seeded_db):
    """带 batch_router 的测试应用：播种 batch-1（user 1，completed+processing）与 batch-2（user 999）。"""
    session_factory = async_sessionmaker(async_engine, class_=AsyncSession, expire_on_commit=False)

    async with session_factory() as session:
        stp_a = STPFile(
            id=9101, user_id=1, object_key="test/batch-a.step", storage_bucket="moldinsight",
            original_filename="batch-a.step", file_size=128, status="completed",
        )
        stp_b = STPFile(
            id=9102, user_id=1, object_key="test/batch-b.step", storage_bucket="moldinsight",
            original_filename="batch-b.step", file_size=128, status="processing",
        )
        stp_c = STPFile(
            id=9103, user_id=999, object_key="test/batch-c.step", storage_bucket="moldinsight",
            original_filename="batch-c.step", file_size=128, status="completed",
        )
        task_a = ProcessingTask(
            task_id="task-batch-a", stp_file_id=9101, batch_id="batch-1",
            task_type="stp_parsing", status="completed", progress=100, parameters={},
        )
        task_b = ProcessingTask(
            task_id="task-batch-b", stp_file_id=9102, batch_id="batch-1",
            task_type="stp_parsing", status="processing", progress=40,
            current_step="生成模具型腔", parameters={},
        )
        task_c = ProcessingTask(
            task_id="task-batch-c", stp_file_id=9103, batch_id="batch-2",
            task_type="stp_parsing", status="failed", progress=40,
            error_message="处理超时", parameters={},
        )
        session.add_all([stp_a, stp_b, stp_c, task_a, task_b, task_c])
        await session.commit()

    test_app = FastAPI()
    test_app.include_router(batch_router, prefix="/api")

    async def override_get_db_session():
        async with session_factory() as session:
            yield session

    test_app.dependency_overrides[get_db_session] = override_get_db_session

    transport = ASGITransport(app=test_app)
    async with AsyncClient(transport=transport, base_url="http://test") as ac:
        yield ac, test_app

    test_app.dependency_overrides.clear()


def _override_user(app: FastAPI, user_id: int):
    app.dependency_overrides[get_current_active_user] = lambda: User(id=user_id, username="tester")


@pytest.mark.asyncio
async def test_batch_status_owner_aggregates_from_pg(batch_client):
    """所有者查询：聚合数字与逐任务字段来自 PG，而非 Redis。"""
    ac, app = batch_client
    _override_user(app, 1)

    resp = await ac.get("/api/batch/batch-1")
    assert resp.status_code == 200
    body = resp.json()

    assert body["batch_id"] == "batch-1"
    assert body["total"] == 2
    assert body["completed"] == 1
    assert body["processing"] == 1
    assert body["failed"] == 0
    assert body["progress_percent"] == 50.0

    by_id = {t["task_id"]: t for t in body["tasks"]}
    assert by_id["task-batch-a"]["status"] == "completed"
    assert by_id["task-batch-a"]["filename"] == "batch-a.step"
    assert by_id["task-batch-b"]["progress"] == 40
    assert by_id["task-batch-b"]["current_step"] == "生成模具型腔"


@pytest.mark.asyncio
async def test_batch_status_non_owner_is_403(batch_client):
    """他人批次必须 403（按 STPFile.user_id 校验，无主不等于公共）。"""
    ac, app = batch_client
    _override_user(app, 2)

    resp = await ac.get("/api/batch/batch-1")
    assert resp.status_code == 403


@pytest.mark.asyncio
async def test_batch_status_unknown_is_404(batch_client):
    ac, app = batch_client
    _override_user(app, 1)

    resp = await ac.get("/api/batch/batch-missing")
    assert resp.status_code == 404


@pytest.mark.asyncio
async def test_batch_status_includes_error_from_pg(batch_client):
    """failed 任务的 error_message 经 PG 返回（此前 error 只存在于 Redis task dict）。"""
    ac, app = batch_client
    _override_user(app, 999)

    resp = await ac.get("/api/batch/batch-2")
    assert resp.status_code == 200
    body = resp.json()
    assert body["failed"] == 1
    assert body["tasks"][0]["error"] == "处理超时"