Pre_Seg_Server/backend/tests/test_tasks.py

from models import ProcessingTask


def test_cancel_task_revokes_celery_and_updates_project(client, db_session, monkeypatch):
    project = client.post("/api/projects", json={
        "name": "Cancelable",
        "video_path": "uploads/1/clip.mp4",
        "status": "parsing",
    }).json()
    task = ProcessingTask(
        task_type="parse_video",
        status="running",
        progress=35,
        message="正在使用 FFmpeg/OpenCV 拆帧",
        project_id=project["id"],
        celery_task_id="celery-1",
        payload={"source_type": "video"},
    )
    db_session.add(task)
    db_session.commit()
    db_session.refresh(task)

    revoked = []
    published = []
    monkeypatch.setattr(
        "routers.tasks.celery_app.control.revoke",
        lambda celery_id, terminate, signal: revoked.append((celery_id, terminate, signal)),
    )
    monkeypatch.setattr("routers.tasks.publish_task_progress_event", lambda event_task: published.append(event_task.status))

    response = client.post(f"/api/tasks/{task.id}/cancel")

    assert response.status_code == 200
    body = response.json()
    assert body["status"] == "cancelled"
    assert body["progress"] == 100
    assert body["message"] == "任务已取消"
    assert body["error"] == "Cancelled by user"
    assert revoked == [("celery-1", True, "SIGTERM")]
    assert published == ["cancelled"]
    assert client.get(f"/api/projects/{project['id']}").json()["status"] == "pending"


def test_retry_task_creates_fresh_parse_task(client, db_session, monkeypatch):
    project = client.post("/api/projects", json={
        "name": "Retryable",
        "video_path": "uploads/2/clip.mp4",
        "source_type": "video",
        "status": "error",
    }).json()
    task = ProcessingTask(
        task_type="parse_video",
        status="failed",
        progress=100,
        message="解析失败",
        error="ffmpeg failed",
        project_id=project["id"],
        payload={"source_type": "video"},
    )
    db_session.add(task)
    db_session.commit()
    db_session.refresh(task)

    class FakeAsyncResult:
        id = "celery-retry"

    queued = []
    published = []
    monkeypatch.setattr("routers.tasks.parse_project_media.delay", lambda task_id: queued.append(task_id) or FakeAsyncResult())
    monkeypatch.setattr("routers.tasks.publish_task_progress_event", lambda event_task: published.append((event_task.id, event_task.status)))

    response = client.post(f"/api/tasks/{task.id}/retry")

    assert response.status_code == 202
    body = response.json()
    assert body["id"] != task.id
    assert body["status"] == "queued"
    assert body["progress"] == 0
    assert body["celery_task_id"] == "celery-retry"
    assert body["payload"]["retry_of"] == task.id
    assert queued == [body["id"]]
    assert published[0] == (body["id"], "queued")
    assert published[-1] == (body["id"], "queued")
    assert client.get(f"/api/projects/{project['id']}").json()["status"] == "parsing"


def test_retry_task_dispatches_propagation_worker_without_media_requirement(client, db_session, monkeypatch):
    project = client.post("/api/projects", json={"name": "Retry Propagation"}).json()
    task = ProcessingTask(
        task_type="propagate_masks",
        status="failed",
        progress=100,
        message="自动传播失败",
        error="model unavailable",
        project_id=project["id"],
        payload={
            "project_id": project["id"],
            "frame_id": 1,
            "steps": [],
        },
    )
    db_session.add(task)
    db_session.commit()
    db_session.refresh(task)

    class FakeAsyncResult:
        id = "celery-propagation-retry"

    queued = []
    monkeypatch.setattr("routers.tasks.propagate_project_masks.delay", lambda task_id: queued.append(task_id) or FakeAsyncResult())
    monkeypatch.setattr("routers.tasks.publish_task_progress_event", lambda event_task: None)

    response = client.post(f"/api/tasks/{task.id}/retry")

    assert response.status_code == 202
    body = response.json()
    assert body["task_type"] == "propagate_masks"
    assert body["celery_task_id"] == "celery-propagation-retry"
    assert queued == [body["id"]]
    assert client.get(f"/api/projects/{project['id']}").json()["status"] == "pending"


def test_task_actions_reject_invalid_states(client, db_session):
    project = client.post("/api/projects", json={
        "name": "Done",
        "video_path": "uploads/3/clip.mp4",
    }).json()
    task = ProcessingTask(
        task_type="parse_video",
        status="success",
        progress=100,
        project_id=project["id"],
        payload={"source_type": "video"},
    )
    db_session.add(task)
    db_session.commit()
    db_session.refresh(task)

    assert client.post(f"/api/tasks/{task.id}/cancel").status_code == 409
    assert client.post(f"/api/tasks/{task.id}/retry").status_code == 409