xiaoxia-saas/tests/unit/test_gpu_lipsync_routes.py

"""GPU Worker 路由单测 — #2009 取消链路.

直接调用路由函数（不经 HTTP 栈），显式注入 svc / _token 以跳过 Depends。
CI 增量映射: gpu_lipsync.py (route) → test_gpu_lipsync_routes.py
"""

from __future__ import annotations

import json
from unittest.mock import MagicMock, patch

import pytest


def _payload(**overrides):
    from app.schemas.gpu_lipsync import GpuWorkerRegisterRequest

    data = {
        "worker_id": "w-1",
        "hostname": "gpu-host",
        "gpu_name": "RTX3060",
        "free_vram_mb": 10000,
        "capabilities": json.dumps({"musetalk": True}),
    }
    data.update(overrides)
    return GpuWorkerRegisterRequest(**data)


def test_register_returns_cancel_task_true_when_cancelled():
    """心跳接口在任务已取消时必须把 cancel_task=True 透传给 Worker."""
    fake_worker = MagicMock()
    fake_worker.worker_id = "w-1"
    fake_worker.hostname = "gpu-host"
    fake_worker.gpu_name = "RTX3060"
    fake_worker.free_vram_mb = 10000
    fake_worker.capabilities = "musetalk"
    fake_svc = MagicMock()
    fake_svc.register_worker.return_value = (fake_worker, True)

    from app.api.routes.gpu_lipsync import register_worker as route

    resp = route(_payload(task_id="task-cancelled"), svc=fake_svc, _token="t")

    assert resp.cancel_task is True
    assert resp.ok is True
    fake_svc.register_worker.assert_called_once()
    kwargs = fake_svc.register_worker.call_args.kwargs
    assert kwargs["task_id"] == "task-cancelled"


def test_register_returns_cancel_task_false_normal():
    """正常心跳 cancel_task=False."""
    fake_worker = MagicMock()
    fake_worker.worker_id = "w-1"
    fake_worker.hostname = "gpu-host"
    fake_worker.gpu_name = "RTX3060"
    fake_worker.free_vram_mb = 10000
    fake_worker.capabilities = "musetalk"
    fake_svc = MagicMock()
    fake_svc.register_worker.return_value = (fake_worker, False)

    from app.api.routes.gpu_lipsync import register_worker as route

    resp = route(_payload(), svc=fake_svc, _token="t")

    assert resp.cancel_task is False


def test_cancel_route_accepts_processing_status():
    """cancel 路由允许 processing 状态（GPU 推理中），不再 400。"""
    fake_job = MagicMock()
    fake_job.status = "cancelled"

    svc = MagicMock()
    svc.cancel_job.return_value = fake_job

    current_user = MagicMock()
    current_user.user.id = "u1"

    from app.api.routes.lipsync import cancel_lipsync_job as route

    result = route("job-1", current_user, svc)

    svc.cancel_job.assert_called_once_with("job-1", "u1")
    assert result.status == "cancelled"


if __name__ == "__main__":
    pytest.main([__file__, "-q"])