xiaoxia-saas/tests/unit/test_duplication_check_worker.py

"""#1661 手动查重 worker task 测试：成功/失败/重试/片段映射/schema 字段。"""

import sys
from pathlib import Path
from unittest.mock import MagicMock, patch

# cv2/numpy 在测试环境不可用，提前 mock
sys.modules.setdefault("cv2", MagicMock())

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


def _get_task(mod):
    """返回 (run_callable, real_task)。

    - celery task 环境：run 是 bound method（self 已绑定），retry 用 patch.object 打桩
    - 原始函数环境：用一个 mock_self 作为 self
    """
    task_obj = mod.process_duplication_check
    real = task_obj._get_current_object() if hasattr(task_obj, "_get_current_object") else task_obj
    if hasattr(real, "run") and hasattr(real, "retry"):
        return real.run, real, True  # bound
    return real, None, False


def _run(mod, record_id, retries=0):
    """执行 task，返回 (result_or_None, raised_exc, mock_self_or_None)。"""
    from celery.exceptions import Retry as CeleryRetry

    func, real_task, bound = _get_task(mod)
    raised = None
    result = None
    if bound:
        mock_retry = MagicMock(side_effect=CeleryRetry("retry"))
        with patch.object(real_task, "retry", mock_retry):
            real_task.request.retries = retries
            real_task.max_retries = 3
            try:
                result = func(record_id)
            except CeleryRetry as e:
                raised = e
        return result, raised, None
    mock_self = MagicMock()
    mock_self.request.retries = retries
    mock_self.max_retries = 3
    mock_self.retry = MagicMock(side_effect=CeleryRetry("retry"))
    try:
        result = func(mock_self, record_id)
    except CeleryRetry as e:
        raised = e
    return result, raised, mock_self


def _make_record(status="pending"):
    from packages.domain.duplication import DuplicationRecord

    record = DuplicationRecord.create(
        user_id="user-1",
        filename="query.mp4",
        file_size=1024,
        storage_key="duplication/abc/query.mp4",
    )
    if status != "pending":
        record.status = status
    return record


def _make_fingerprint():
    from video_processing.dedup import FingerprintChunk, VideoFingerprint

    chunks = [
        FingerprintChunk(start_time_ms=0, end_time_ms=2000, phash_binary="0" * 16, color_histogram=[], frame_count=1),
        FingerprintChunk(
            start_time_ms=2000, end_time_ms=4000, phash_binary="1" * 16, color_histogram=[], frame_count=1
        ),
    ]
    return VideoFingerprint(
        md5="qmd5",
        keyframe_phashes=[c.phash_binary for c in chunks],
        color_histograms=[],
        duration=10000.0,
        resolution=(720, 1280),
        chunks=chunks,
    )


def _patch_common(record, storage=None, dedup=None, session=None):
    from worker_app.tasks import duplication_check as mod

    fake_repo = MagicMock()
    fake_repo.get.return_value = record
    return [
        patch.object(mod, "SessionLocal", return_value=session or MagicMock()),
        patch.object(mod, "SQLAlchemyDuplicationRecordRepository", return_value=fake_repo),
        patch.object(mod, "get_storage_service", return_value=storage or MagicMock()),
        patch.object(mod, "VideoDeduplicator", return_value=dedup or MagicMock()),
    ], fake_repo


class TestProcessDuplicationCheckSuccess:
    def test_success_flow_updates_record(self):
        from worker_app.tasks import duplication_check as mod

        record = _make_record()
        fake_session = MagicMock()
        fake_storage = MagicMock()
        fake_dedup = MagicMock()
        fake_dedup.compute_fingerprint.return_value = _make_fingerprint()
        fake_dedup.compute_duplicate_rate.return_value = {
            "duplicate_rate": 42.5,
            "visual_similarity": 0.83,
            "match_count": 1,
        }
        patches, fake_repo = _patch_common(record, storage=fake_storage, dedup=fake_dedup, session=fake_session)
        patches.append(patch.object(mod, "_build_domain_segments", return_value=(["SEG"], 1)))
        for p in patches:
            p.start()
        try:
            result, raised, _ = _run(mod, record.id)
        finally:
            for p in patches:
                p.stop()

        assert raised is None
        assert result["ok"] is True
        assert result["status"] == "completed"
        assert result["duplicate_rate"] == 42.5
        assert result["visual_similarity"] == 0.83
        assert result["match_count"] == 1
        assert result["segments"] == 1

        assert record.status == "completed"
        assert record.duplicate_rate == 42.5
        assert record.visual_similarity == 0.83
        assert record.match_count == 1
        assert record.duplicate_count == 1
        assert record.segments == ["SEG"]

        fake_storage.download_file.assert_called_once()
        fake_dedup.compute_fingerprint.assert_called_once()
        _, kwargs = fake_dedup.compute_duplicate_rate.call_args
        assert kwargs["scope"] == "user"
        assert kwargs["user_id"] == "user-1"
        assert kwargs["current_video_id"] is None
        assert fake_repo.update.call_count >= 2
        fake_session.commit.assert_called()
        fake_session.close.assert_called()

    def test_already_completed_is_skipped(self):
        from worker_app.tasks import duplication_check as mod

        record = _make_record(status="completed")
        patches, fake_repo = _patch_common(record)
        for p in patches:
            p.start()
        try:
            result, raised, _ = _run(mod, record.id)
        finally:
            for p in patches:
                p.stop()
        assert raised is None
        assert result.get("skipped") is True
        fake_repo.update.assert_not_called()


class TestProcessDuplicationCheckFailure:
    def test_record_not_found_raises(self):
        from worker_app.tasks import duplication_check as mod

        fake_repo = MagicMock()
        fake_repo.get.return_value = None
        patches = [
            patch.object(mod, "SessionLocal", return_value=MagicMock()),
            patch.object(mod, "SQLAlchemyDuplicationRecordRepository", return_value=fake_repo),
            patch.object(mod, "get_storage_service", return_value=MagicMock()),
        ]
        for p in patches:
            p.start()
        try:
            _result, raised, _ = _run(mod, "nope", retries=0)
        finally:
            for p in patches:
                p.stop()
        # 找不到记录触发异常 → retry（第一次）
        assert raised is not None

    def test_download_failure_retries_then_marks_failed(self):
        from worker_app.tasks import duplication_check as mod

        # 第一次失败（retries=0）：保持 pending
        record = _make_record()
        fake_storage = MagicMock()
        fake_storage.download_file.side_effect = RuntimeError("oss network down")
        patches, _ = _patch_common(record, storage=fake_storage)
        for p in patches:
            p.start()
        try:
            _, raised, _ = _run(mod, record.id, retries=0)
        finally:
            for p in patches:
                p.stop()
        assert raised is not None
        assert record.status == "processing", "首次失败不应标记 failed（已进入 processing 等待重试）"

        # 最后一次（retries==max_retries=3）：标记 failed
        record2 = _make_record()
        patches2, fake_repo2 = _patch_common(record2, storage=fake_storage)
        for p in patches2:
            p.start()
        try:
            _run(mod, record2.id, retries=3)
        finally:
            for p in patches2:
                p.stop()
        assert record2.status == "failed"
        assert "查重失败" in record2.error_message
        fake_repo2.update.assert_called()

    def test_temp_dir_cleaned_after_failure(self):
        import os
        import tempfile

        from worker_app.tasks import duplication_check as mod

        record = _make_record()
        fake_storage = MagicMock()
        fake_storage.download_file.side_effect = RuntimeError("boom")

        created_dirs = []
        real_mkdtemp = tempfile.mkdtemp

        def fake_mkdtemp(prefix=None):
            d = real_mkdtemp(prefix=prefix)
            created_dirs.append(d)
            return d

        patches, _ = _patch_common(record, storage=fake_storage)
        patches.append(patch.object(mod.tempfile, "mkdtemp", fake_mkdtemp))
        for p in patches:
            p.start()
        try:
            _run(mod, record.id, retries=0)
        finally:
            for p in patches:
                p.stop()

        assert created_dirs, "mkdtemp should have been called"
        assert not os.path.isdir(created_dirs[0]), "temp dir should be removed in finally"


class TestBuildDomainSegments:
    def test_maps_worker_segments_to_domain_with_seconds_and_percent(self):
        from video_processing.dedup import DuplicateSegment as WorkerSegment
        from worker_app.tasks import duplication_check as mod

        fingerprint = _make_fingerprint()

        from packages.domain import GeneratedVideo

        existing = GeneratedVideo(
            id="vid-1",
            project_id="proj-1",
            generation_task_id="t1",
            name="成片A",
            file_url="oss://x",
            file_size=1,
            duration=10.0,
            width=720,
            height=1280,
            fps=30.0,
            video_fingerprint={"md5": "x"},
        )
        fake_video_repo = MagicMock()
        fake_video_repo.list_by_user.return_value = [existing]

        fake_dedup = MagicMock()
        fake_dedup._get_existing_chunks.return_value = [
            {"phash_binary": "0" * 16, "start_time_ms": 0, "end_time_ms": 2000, "color_histogram": []},
        ]
        worker_seg = WorkerSegment(
            query_start_ms=1000,
            query_end_ms=3000,
            target_start_ms=5000,
            target_end_ms=7000,
            avg_distance=6.0,
        )

        with (
            patch.object(mod, "SQLAlchemyGeneratedVideoRepository", return_value=fake_video_repo),
            patch.object(mod, "find_duplicate_segments", return_value=[worker_seg]),
        ):
            segments, dup_count = mod._build_domain_segments(fingerprint, MagicMock(), fake_dedup, "user-1")

        assert dup_count == 1
        assert len(segments) == 1
        seg = segments[0]
        assert seg.source_start == 1.0
        assert seg.source_end == 3.0
        assert seg.matched_start == 5.0
        assert seg.matched_end == 7.0
        assert seg.matched_video_id == "vid-1"
        assert seg.matched_video_name == "成片A"
        assert abs(seg.similarity - 90.6) < 0.2

    def test_skips_videos_without_chunks(self):
        from worker_app.tasks import duplication_check as mod

        fingerprint = _make_fingerprint()
        from packages.domain import GeneratedVideo

        existing = GeneratedVideo(
            id="vid-2",
            project_id="p",
            generation_task_id="t",
            name="老视频",
            file_url="oss://x",
            file_size=1,
            duration=5.0,
            width=720,
            height=1280,
            fps=30.0,
            video_fingerprint={"md5": "old"},
        )
        fake_video_repo = MagicMock()
        fake_video_repo.list_by_user.return_value = [existing]
        fake_dedup = MagicMock()
        fake_dedup._get_existing_chunks.return_value = []

        with patch.object(mod, "SQLAlchemyGeneratedVideoRepository", return_value=fake_video_repo):
            segments, dup_count = mod._build_domain_segments(fingerprint, MagicMock(), fake_dedup, "u")
        assert segments == []
        assert dup_count == 0


class TestDuplicationSchemaAndDomainNewFields:
    def test_record_response_includes_new_fields(self):
        from app.schemas.duplication import DuplicationRecordResponse

        resp = DuplicationRecordResponse(
            id="r1",
            filename="f.mp4",
            file_size=1,
            status="completed",
            duplicate_rate=10.0,
            duplicate_count=1,
            visual_similarity=0.5,
            match_count=2,
            created_at="2026-09-04T00:00:00",
            updated_at="2026-09-04T00:00:00",
        )
        assert resp.visual_similarity == 0.5
        assert resp.match_count == 2

    def test_record_response_new_fields_default_none(self):
        from app.schemas.duplication import DuplicationRecordResponse

        resp = DuplicationRecordResponse(id="r1", filename="f.mp4", file_size=1, created_at="x", updated_at="y")
        assert resp.visual_similarity is None
        assert resp.match_count is None

    def test_domain_mark_completed_accepts_new_fields(self):
        record = _make_record()
        record.mark_completed(33.0, 2, [], visual_similarity=0.77, match_count=3)
        assert record.status == "completed"
        assert record.visual_similarity == 0.77
        assert record.match_count == 3

    def test_reset_for_retry_clears_new_fields(self):
        record = _make_record()
        record.mark_completed(10.0, 1, [], visual_similarity=0.5, match_count=1)
        record.status = "failed"
        record.reset_for_retry()
        assert record.status == "pending"
        assert record.visual_similarity is None
        assert record.match_count is None