xiaoxia-saas/tests/unit/test_generation_task_orphan_cleanup.py

"""GenerationTaskRepository - cleanup_stale_running 孤儿任务清理单元测试。"""

import sys
from datetime import datetime, timedelta, timezone
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "apps" / "api"))

from sqlalchemy import create_engine, text
from sqlalchemy.orm import sessionmaker

from packages.adapters.sqlalchemy_impl.generation_task_repository import (
    SQLAlchemyGenerationTaskRepository,
)
from packages.adapters.sqlalchemy_impl.models import Base
from packages.domain import GenerationTask, GenerationTaskStatus


def _repository():
    engine = create_engine("sqlite:///:memory:")
    Base.metadata.create_all(engine)
    session = sessionmaker(bind=engine)()
    return SQLAlchemyGenerationTaskRepository(session), session, engine


def _make_task(**kwargs) -> GenerationTask:
    defaults = dict(
        project_id="proj-1",
        asset_library_id="lib-1",
        created_by_user_id="user-1",
    )
    defaults.update(kwargs)
    return GenerationTask.create(**defaults)


def test_cleanup_stale_running_no_tasks_returns_zero():
    """没有任务时返回0。"""
    repo, _, _ = _repository()
    count = repo.cleanup_stale_running(timeout_minutes=10)
    assert count == 0


def test_cleanup_stale_running_recent_running_not_cleaned():
    """刚启动的 running 任务不清理。"""
    repo, _, _ = _repository()
    task = _make_task()
    repo.create(task)
    task.mark_processing()
    repo.update(task)

    count = repo.cleanup_stale_running(timeout_minutes=10)
    assert count == 0
    assert repo.get(task.id).status == GenerationTaskStatus.RUNNING


def test_cleanup_stale_running_old_running_marked_failed():
    """超时的 running 任务被标记为 failed。"""
    repo, _, engine = _repository()
    task = _make_task()
    repo.create(task)
    task.mark_processing()
    repo.update(task)

    # 手动把 updated_at 改到 30 分钟前
    with engine.connect() as conn:
        conn.execute(
            text("UPDATE generation_tasks SET updated_at = :ts WHERE id = :id"),
            {"ts": datetime.now(timezone.utc) - timedelta(minutes=30), "id": task.id},
        )
        conn.commit()

    count = repo.cleanup_stale_running(timeout_minutes=10)
    assert count == 1

    saved = repo.get(task.id)
    assert saved.status == GenerationTaskStatus.FAILED
    assert "任务执行中断" in saved.error_message
    assert saved.error_info.get("error_type") == "WorkerInterrupted"
    assert saved.completed_at is not None


def test_cleanup_stale_running_pending_not_touched():
    """pending 状态即使超时也不清理。"""
    repo, _, engine = _repository()
    task = _make_task()
    repo.create(task)

    with engine.connect() as conn:
        conn.execute(
            text("UPDATE generation_tasks SET updated_at = :ts WHERE id = :id"),
            {"ts": datetime.now(timezone.utc) - timedelta(hours=1), "id": task.id},
        )
        conn.commit()

    count = repo.cleanup_stale_running(timeout_minutes=10)
    assert count == 0
    assert repo.get(task.id).status == GenerationTaskStatus.PENDING


def test_cleanup_stale_running_multiple_orphans():
    """多个超时 running 任务全部清理。"""
    repo, _, engine = _repository()

    tasks = []
    for i in range(3):
        t = _make_task(project_id=f"proj-{i}")
        repo.create(t)
        t.mark_processing()
        repo.update(t)
        tasks.append(t)

    # 前两个超时，第三个是新的
    with engine.connect() as conn:
        for t in tasks[:2]:
            conn.execute(
                text("UPDATE generation_tasks SET updated_at = :ts WHERE id = :id"),
                {"ts": datetime.now(timezone.utc) - timedelta(minutes=20), "id": t.id},
            )
        conn.commit()

    count = repo.cleanup_stale_running(timeout_minutes=10)
    assert count == 2

    assert repo.get(tasks[0].id).status == GenerationTaskStatus.FAILED
    assert repo.get(tasks[1].id).status == GenerationTaskStatus.FAILED
    assert repo.get(tasks[2].id).status == GenerationTaskStatus.RUNNING