xiaoxia-saas/tests/unit/test_ingest_hevc_transcode_task.py

"""ingest_asset 任务中 HEVC 转码主流程的任务级单元测试。

通过 mock subprocess / repository / OSS，验证：
- 转码成功 + 方向校验通过 → storage_key 改写为 *_h264
- 方向校验失败（竖屏转出横屏）→ 降级原文件，storage_key 不变，error 日志
- ffmpeg 非零退出 → 降级原文件
- 非 HEVC 编码 → 不触发转码
"""

from __future__ import annotations

import sys
import tempfile
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import MagicMock, patch

# 在 import worker_app 模块前 mock 掉数据库连接和 celery（同 test_ingest_validation.py）
_mock_db_module = MagicMock()
_mock_db_module.SessionLocal = MagicMock()
sys.modules["worker_app.db"] = _mock_db_module
sys.modules["worker_app.core.config"] = MagicMock()

_mock_celery_module = MagicMock()


def _passthrough_decorator(*args, **kwargs):
    if len(args) == 1 and callable(args[0]):
        return args[0]
    return lambda f: f


_mock_celery_module.celery_app.task = MagicMock(side_effect=_passthrough_decorator)
sys.modules["worker_app.celery_app"] = _mock_celery_module

# mock video_processing 子模块（主流程会 import 它们）
_oss_helpers_mock = MagicMock()
_thumbnail_mock = MagicMock()
sys.modules["video_processing.oss_helpers"] = _oss_helpers_mock
sys.modules["video_processing.thumbnail_generator"] = _thumbnail_mock

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

import pytest  # noqa: E402
from worker_app.tasks import ingest as ingest_mod  # noqa: E402


class _FakeJobRepo:
    def __init__(self, db):
        self.initial_job = SimpleNamespace(
            id="job-1",
            project_id="proj-1",
            library_id="lib-1",
            storage_key="uploads/proj/IMG_2281.MOV",
            file_hash="hash-1",
            status=None,
            error_message=None,
            result_asset_id=None,
            updated_at=None,
        )
        self.updated_job = None

    def get(self, job_id):
        return self.initial_job

    def update(self, job):
        # 生产代码在同一 job 对象上原地修改属性后传入 update；
        # 这里捕获引用，断言时读最终状态
        self.updated_job = job

    @property
    def final_job(self):
        return self.updated_job or self.initial_job


class _FakeAssetRepo:
    def __init__(self, db):
        self.created = None

    def create(self, asset):
        self.created = asset


def _video_metadata(codec="hevc"):
    return {
        "codec": codec,
        "width": 1920,
        "height": 1080,
        "duration": 10.0,
        "size_bytes": 5 * 1024 * 1024,
    }


@pytest.fixture
def task_env(tmp_path):
    """统一构造 ingest_asset 主流程的 mock 环境。返回控制句柄。

    测试中用 mocks = _start_patches(control) 启动，断言必须用
    mocks["upload"] 等 start() 返回的 mock；不能在 stop() 后读模块
    属性（stop 后属性恢复为原 auto-mock，调用记录为 0）。
    """
    local_file = tmp_path / "local_hevc.MOV"
    local_file.write_bytes(b"fake-hevc-source")
    tc_out = tmp_path / "transcode_out_h264.mp4"

    control = {
        "rotation_source": 90,  # 源文件 rotation；None=横屏无 side data
        "transcode_rc": 0,
        "transcode_produces_file": True,
        "validate_ok": True,
        "upload_url": "https://oss.example.com/x_h264.MOV",
        "codec": "hevc",
        "tc_out": tc_out,
        "local_file": local_file,
    }

    def fake_probe_rotation(path):
        if Path(path).name == tc_out.name:
            return None  # 产物无 side data
        return control["rotation_source"]

    def fake_probe_dimensions(path):
        if Path(path).name == tc_out.name:
            return (1080, 1920) if control["validate_ok"] else (1920, 1080)
        return (1920, 1080)

    control["subprocess_calls"] = []

    def fake_subprocess_run(cmd, **kwargs):
        control["subprocess_calls"].append(list(cmd[:3]))
        if cmd and cmd[0] == "ffmpeg" and "libx264" in cmd:
            if control["transcode_rc"] != 0:
                return SimpleNamespace(returncode=control["transcode_rc"], stderr="boom")
            if control["transcode_produces_file"]:
                Path(cmd[-1]).write_bytes(b"fake-h264-output")
            return SimpleNamespace(returncode=0, stderr="")
        return SimpleNamespace(returncode=0, stdout="", stderr="")

    def fake_ntf(*args, **kwargs):
        mock_file = MagicMock()
        mock_file.name = str(tc_out) if kwargs.get("suffix") == "_h264.mp4" else str(local_file)
        mock_file.close = MagicMock()
        # with ... as tmp: 让 __enter__ 返回自身，tmp.name 才是上面设置的路径
        mock_file.__enter__.return_value = mock_file
        mock_file.__exit__.return_value = False
        return mock_file

    job_repo = _FakeJobRepo(db=None)
    asset_repo = _FakeAssetRepo(db=None)

    control["patchers"] = {
        "session": patch.object(ingest_mod, "SessionLocal", return_value=MagicMock()),
        "job_repo": patch.object(ingest_mod, "SQLAlchemyIngestJobRepository", return_value=job_repo),
        "asset_repo": patch.object(ingest_mod, "SQLAlchemyAssetRepository", return_value=asset_repo),
        "download": patch.object(ingest_mod, "download_asset", return_value=True),
        "upload": patch.object(
            sys.modules["video_processing.oss_helpers"],
            "upload_to_oss",
            return_value=control["upload_url"],
        ),
        "metadata": patch.object(
            ingest_mod,
            "extract_media_metadata",
            side_effect=lambda path, mt: (
                (_video_metadata("h264"), True)
                if Path(path).name == tc_out.name
                else (_video_metadata(control["codec"]), True)
            ),
        ),
        "rotation": patch.object(ingest_mod, "probe_rotation", side_effect=fake_probe_rotation),
        "dimensions": patch.object(ingest_mod, "probe_dimensions", side_effect=fake_probe_dimensions),
        "validate": patch.object(
            ingest_mod,
            "validate_transcode_output",
            side_effect=lambda p, portrait: control["validate_ok"],
        ),
        "subprocess": patch.object(ingest_mod.subprocess, "run", side_effect=fake_subprocess_run),
        "ntf": patch.object(tempfile, "NamedTemporaryFile", side_effect=fake_ntf),
        # 缩略图生成跳过
        "thumb": patch(
            "video_processing.thumbnail_generator.extract_first_frame",
            side_effect=RuntimeError("skip thumb"),
        ),
    }

    control["job_repo"] = job_repo
    control["asset_repo"] = asset_repo
    return control


def _start_patches(control):
    """启动全部 patcher，返回具名 mock dict（调用记录都在这些 mock 上）。"""
    return {name: p.start() for name, p in control["patchers"].items()}


def _stop_patches(control):
    for p in control["patchers"].values():
        p.stop()


class TestIngestHEVCTranscodeFlow:
    def test_success_rewrites_storage_key(self, task_env):
        """竖屏 HEVC 转码+校验通过 → storage_key 改写为 *_h264.MOV，asset READY 入库。"""
        mocks = _start_patches(task_env)
        try:
            result = ingest_mod.ingest_asset("job-1")
        finally:
            _stop_patches(task_env)

        assert result["status"] == "completed"
        assert task_env["job_repo"].final_job.storage_key == "uploads/proj/IMG_2281_h264.MOV"
        assert task_env["asset_repo"].created is not None
        # 转码产物上传 OSS 恰好一次，且上传的是 *_h264.MOV 新 key
        mocks["upload"].assert_called_once()
        uploaded_path, uploaded_key = mocks["upload"].call_args.args
        assert uploaded_key == "uploads/proj/IMG_2281_h264.MOV"
        assert str(uploaded_path).endswith("_h264.mp4")

    def test_validation_failure_keeps_original_file(self, task_env):
        """竖屏转出横屏（校验失败）→ 降级原文件，storage_key 不变，打 error 日志。"""
        task_env["validate_ok"] = False
        mocks = _start_patches(task_env)
        error_mock = MagicMock()
        try:
            with patch.object(ingest_mod.logger, "error", error_mock):
                result = ingest_mod.ingest_asset("job-1")
        finally:
            _stop_patches(task_env)

        assert result["status"] == "completed"
        assert task_env["job_repo"].final_job.storage_key == "uploads/proj/IMG_2281.MOV"
        assert error_mock.called
        assert "方向/维度校验失败" in error_mock.call_args[0][0]
        # 校验失败：转码产物不得上传 OSS，杜绝横屏文件覆盖
        mocks["upload"].assert_not_called()

    def test_ffmpeg_nonzero_keeps_original(self, task_env):
        """ffmpeg 返回非零 → 降级原文件，storage_key 不变。"""
        task_env["transcode_rc"] = 1
        mocks = _start_patches(task_env)
        try:
            result = ingest_mod.ingest_asset("job-1")
        finally:
            _stop_patches(task_env)

        assert result["status"] == "completed"
        assert task_env["job_repo"].final_job.storage_key == "uploads/proj/IMG_2281.MOV"
        mocks["upload"].assert_not_called()

    def test_non_hevc_no_transcode(self, task_env):
        """非 HEVC 编码（h264）→ 不触发 ffmpeg 转码。"""
        task_env["codec"] = "h264"
        task_env["rotation_source"] = None
        mocks = _start_patches(task_env)
        try:
            result = ingest_mod.ingest_asset("job-1")
        finally:
            _stop_patches(task_env)

        assert result["status"] == "completed"
        assert task_env["job_repo"].final_job.storage_key == "uploads/proj/IMG_2281.MOV"
        mocks["upload"].assert_not_called()
        # 所有 subprocess 调用都不应是 ffmpeg 转码
        for call in mocks["subprocess"].call_args_list:
            cmd = call.args[0] if call.args else call.kwargs.get("cmd", [])
            assert "libx264" not in cmd