xiaoxia-saas/tests/unit/test_dedup_engine.py

"""查重引擎单元测试。

覆盖：
- hamming_distance() 汉明距离计算（XOR bit 计数）
- compute_phash() 感知哈希算法（需真实 cv2，无则跳过）
- compute_color_histogram() 颜色直方图（需真实 cv2，无则跳过）
- VideoDeduplicator.check_duplicate() 相似度判定逻辑
"""

from __future__ import annotations

import sys
from unittest.mock import MagicMock

# ---------------------------------------------------------------------------
# 保存 sys.modules 原始状态，测试结束后恢复，避免污染其他测试文件
# ---------------------------------------------------------------------------
_ORIGINAL_MODULES = dict(sys.modules)
_MOCKED_MODULE_NAMES: list[str] = []


def _mock_if_absent(name: str, mock_obj=None):
    """仅在模块不在 sys.modules 中时注入 mock，并记录以便清理。"""
    if name not in sys.modules:
        sys.modules[name] = mock_obj if mock_obj is not None else MagicMock()
        _MOCKED_MODULE_NAMES.append(name)


# Mock heavy deps before importing dedup module
_mock_if_absent("ffmpeg")

# Mock worker_app (celery) and its submodules
for mod_name in ["worker_app", "worker_app.celery_app", "worker_app.db"]:
    _mock_if_absent(mod_name)
if "worker_app.celery_app" in sys.modules and isinstance(sys.modules["worker_app.celery_app"], MagicMock):
    sys.modules["worker_app.celery_app"].celery_app = MagicMock()
if "worker_app.db" in sys.modules and isinstance(sys.modules["worker_app.db"], MagicMock):
    sys.modules["worker_app.db"].SessionLocal = MagicMock()

# Mock celery.Task base class
_mock_if_absent("celery", MagicMock())
if "celery" in sys.modules and isinstance(sys.modules["celery"], MagicMock):
    sys.modules["celery"].Task = object

# Mock packages.shared.storage
_mock_if_absent("packages.shared")
_mock_if_absent("packages.shared.storage")

# Mock packages.adapters.sqlalchemy_impl.generated_video_repository
_mock_if_absent("packages.adapters.sqlalchemy_impl.generated_video_repository")

# Check if cv2 is available as a real module (not mocked)
_HAS_CV2 = False
try:
    import cv2 as _cv2
    if not isinstance(_cv2, MagicMock):
        _HAS_CV2 = True
except (ImportError, ModuleNotFoundError):
    pass

import numpy as np  # noqa: E402
import pytest  # noqa: E402

# Mock cv2 if not available (so dedup module can import)
if not _HAS_CV2:
    _mock_if_absent("cv2")

from apps.worker.video_processing.dedup import (  # noqa: E402
    VideoDeduplicator,
    VideoFingerprint,
    compute_color_histogram,
    compute_phash,
    hamming_distance,
)

# ---------------------------------------------------------------------------
# dedup 模块已导入完成，立即恢复 worker_app 真实包，避免污染后续测试文件
# ---------------------------------------------------------------------------
for _name in ["worker_app", "worker_app.celery_app", "worker_app.db"]:
    if _name in _MOCKED_MODULE_NAMES:
        sys.modules.pop(_name, None)
        _MOCKED_MODULE_NAMES.remove(_name)


@pytest.fixture(autouse=True, scope="session")
def _cleanup_mocks():
    """测试结束后恢复 sys.modules，防止 mock 污染其他测试文件。"""
    yield
    # 移除本次新增的 mock 模块
    for name in _MOCKED_MODULE_NAMES:
        sys.modules.pop(name, None)
    # 恢复被覆盖的模块
    for name, mod in _ORIGINAL_MODULES.items():
        if sys.modules.get(name) is not mod:
            sys.modules[name] = mod


class TestHammingDistance:
    """hamming_distance() 测试。

    实现使用 XOR + bit 计数：bin(h1 ^ h2).count("1")。
    空字符串会触发 ValueError（int("", 16) 失败），属于边界行为。
    """

    def test_identical_hashes_zero_distance(self):
        assert hamming_distance("abcdef01", "abcdef01") == 0

    def test_completely_different_bytes(self):
        # 0x00 XOR 0xFF = 0xFF → 8 bits
        assert hamming_distance("00", "ff") == 8

    def test_single_bit_difference(self):
        # 0x0 XOR 0x1 = 0x1 → 1 bit
        assert hamming_distance("0", "1") == 1

    def test_unequal_length_leading_zeros(self):
        # int("abc", 16) == int("0abc", 16) → XOR = 0 → 0 bits
        dist = hamming_distance("abc", "0abc")
        assert dist == 0

    def test_unequal_length_with_leading_zeros_ff(self):
        # int("ff", 16) == int("00ff", 16) → XOR = 0 → 0 bits
        dist = hamming_distance("ff", "00ff")
        assert dist == 0

    def test_all_bits_different_64bit(self):
        # 16 hex chars = 64 bits, all different → 64
        dist = hamming_distance("0000000000000000", "ffffffffffffffff")
        assert dist == 64

    def test_partial_difference(self):
        # 0x0F = 00001111, 0xF0 = 11110000 → XOR = 0xFF → 8 bits
        assert hamming_distance("0f", "f0") == 8

    def test_one_bit_in_second_byte(self):
        # 0x0000 XOR 0x0001 = 0x0001 → 1 bit
        assert hamming_distance("0000", "0001") == 1


@pytest.mark.skipif(not _HAS_CV2, reason="需要真实 cv2 模块")
class TestComputePhash:
    """compute_phash() 测试（需真实 cv2）。"""

    def test_returns_hex_string(self):
        image = np.random.randint(0, 256, (64, 64, 3), dtype=np.uint8)
        result = compute_phash(image)
        assert isinstance(result, str)
        int(result, 16)  # 不应抛出异常

    def test_identical_images_same_hash(self):
        image = np.full((64, 64, 3), 128, dtype=np.uint8)
        hash1 = compute_phash(image)
        hash2 = compute_phash(image)
        assert hash1 == hash2

    def test_different_images_different_hash(self):
        img1 = np.zeros((64, 64, 3), dtype=np.uint8)
        img2 = np.full((64, 64, 3), 255, dtype=np.uint8)
        hash1 = compute_phash(img1)
        hash2 = compute_phash(img2)
        assert hash1 != hash2

    def test_custom_hash_size(self):
        image = np.random.randint(0, 256, (64, 64, 3), dtype=np.uint8)
        result = compute_phash(image, hash_size=16)
        assert isinstance(result, str)
        int(result, 16)


@pytest.mark.skipif(not _HAS_CV2, reason="需要真实 cv2 模块")
class TestComputeColorHistogram:
    """compute_color_histogram() 测试（需真实 cv2）。"""

    def test_returns_correct_length(self):
        image = np.random.randint(0, 256, (64, 64, 3), dtype=np.uint8)
        hist = compute_color_histogram(image, bins=32)
        assert len(hist) == 96  # 3 channels × 32 bins

    def test_custom_bins(self):
        image = np.random.randint(0, 256, (64, 64, 3), dtype=np.uint8)
        hist = compute_color_histogram(image, bins=16)
        assert len(hist) == 48  # 3 channels × 16 bins

    def test_normalized_values(self):
        image = np.random.randint(0, 256, (64, 64, 3), dtype=np.uint8)
        hist = compute_color_histogram(image)
        for v in hist:
            assert 0.0 <= v <= 1.0 + 1e-6

    def test_identical_images_same_histogram(self):
        image = np.full((64, 64, 3), 100, dtype=np.uint8)
        hist1 = compute_color_histogram(image)
        hist2 = compute_color_histogram(image)
        assert hist1 == hist2


class TestVideoDeduplicatorCheckDuplicate:
    """VideoDeduplicator.check_duplicate() 测试。

    当前实现仅使用 MD5 精确匹配和 pHash 距离判定，
    不包含颜色直方图相似度计算。
    """

    @pytest.fixture
    def deduplicator(self):
        return VideoDeduplicator()

    @pytest.fixture
    def mock_session(self):
        return MagicMock()

    def _make_existing_video(self, video_id, md5, phashes=None):
        """创建模拟已有视频的 mock 对象。"""
        video = MagicMock()
        video.id = video_id
        video.video_fingerprint = {
            "md5": md5,
            "keyframe_phashes": phashes or [],
            "color_histograms": [],
        }
        return video

    def _patch_repo(self, mock_repo):
        """Patch SQLAlchemyGeneratedVideoRepository。"""
        import apps.worker.video_processing.dedup as dedup_module
        original = dedup_module.SQLAlchemyGeneratedVideoRepository
        dedup_module.SQLAlchemyGeneratedVideoRepository = MagicMock(return_value=mock_repo)
        return original, dedup_module

    def _restore_repo(self, dedup_module, original):
        dedup_module.SQLAlchemyGeneratedVideoRepository = original

    def test_exact_md5_match(self, deduplicator, mock_session):
        """MD5 完全匹配应返回 similarity=1.0。"""
        existing = self._make_existing_video("vid-1", "abc123")
        mock_repo = MagicMock()
        mock_repo.list_by_project.return_value = [existing]

        fingerprint = VideoFingerprint(
            md5="abc123",
            keyframe_phashes=["ff"],
            color_histograms=[],
            duration=10.0,
            resolution=(1280, 720),
        )

        orig, mod = self._patch_repo(mock_repo)
        try:
            result = deduplicator.check_duplicate(fingerprint, "proj-1", mock_session)
            assert result is not None
            assert result["duplicate"] is True
            assert result["similarity"] == 1.0
            assert result["reason"] == "exact_md5_match"
        finally:
            self._restore_repo(mod, orig)

    def test_phash_similar_match(self, deduplicator, mock_session):
        """pHash 距离 < 阈值时应判定为重复。"""
        existing = self._make_existing_video("vid-1", "different_md5", phashes=["abcdef01"])
        mock_repo = MagicMock()
        mock_repo.list_by_project.return_value = [existing]

        fingerprint = VideoFingerprint(
            md5="different_md5_new",
            keyframe_phashes=["abcdef01"],  # 完全相同，距离=0
            color_histograms=[],
            duration=10.0,
            resolution=(1280, 720),
        )

        orig, mod = self._patch_repo(mock_repo)
        try:
            result = deduplicator.check_duplicate(fingerprint, "proj-1", mock_session)
            assert result is not None
            assert result["duplicate"] is True
            assert result["similarity"] == 1.0  # distance=0 → 1.0
            assert result["reason"] == "phash_similar"
        finally:
            self._restore_repo(mod, orig)

    def test_no_match_returns_none(self, deduplicator, mock_session):
        """pHash 平均距离 >= PHASH_THRESHOLD(10) 时应返回 None。"""
        # 使用 16 字符 phash（64 bit），全部不同 → 距离=64 >= 10
        existing = self._make_existing_video("vid-1", "md5_a", phashes=["0000000000000000"])
        mock_repo = MagicMock()
        mock_repo.list_by_project.return_value = [existing]

        fingerprint = VideoFingerprint(
            md5="md5_b",
            keyframe_phashes=["ffffffffffffffff"],  # 64 bits 全不同 → 距离=64
            color_histograms=[],
            duration=10.0,
            resolution=(1280, 720),
        )

        orig, mod = self._patch_repo(mock_repo)
        try:
            result = deduplicator.check_duplicate(fingerprint, "proj-1", mock_session)
            assert result is None
        finally:
            self._restore_repo(mod, orig)

    def test_empty_project_returns_none(self, deduplicator, mock_session):
        """项目中没有视频时应返回 None。"""
        mock_repo = MagicMock()
        mock_repo.list_by_project.return_value = []

        fingerprint = VideoFingerprint(
            md5="abc",
            keyframe_phashes=["ff"],
            color_histograms=[],
            duration=10.0,
            resolution=(1280, 720),
        )

        orig, mod = self._patch_repo(mock_repo)
        try:
            result = deduplicator.check_duplicate(fingerprint, "proj-1", mock_session)
            assert result is None
        finally:
            self._restore_repo(mod, orig)

    def test_skip_videos_without_fingerprint(self, deduplicator, mock_session):
        """没有指纹的视频应被跳过。"""
        existing = MagicMock()
        existing.id = "vid-1"
        existing.video_fingerprint = None

        mock_repo = MagicMock()
        mock_repo.list_by_project.return_value = [existing]

        fingerprint = VideoFingerprint(
            md5="abc",
            keyframe_phashes=["ff"],
            color_histograms=[],
            duration=10.0,
            resolution=(1280, 720),
        )

        orig, mod = self._patch_repo(mock_repo)
        try:
            result = deduplicator.check_duplicate(fingerprint, "proj-1", mock_session)
            assert result is None
        finally:
            self._restore_repo(mod, orig)

    def test_first_match_returned(self, deduplicator, mock_session):
        """返回第一个通过阈值的匹配（非最优匹配）。"""
        # vid-1: 距离=2 bits（0x03 XOR 0x01 = 0x02 → 1 bit），通过阈值
        vid1 = self._make_existing_video("vid-1", "md5_1", phashes=["0000000000000003"])
        # vid-2: 距离=0 bits（完全匹配）
        vid2 = self._make_existing_video("vid-2", "md5_2", phashes=["0000000000000001"])

        mock_repo = MagicMock()
        mock_repo.list_by_project.return_value = [vid1, vid2]

        fingerprint = VideoFingerprint(
            md5="md5_new",
            keyframe_phashes=["0000000000000001"],
            color_histograms=[],
            duration=10.0,
            resolution=(1280, 720),
        )

        orig, mod = self._patch_repo(mock_repo)
        try:
            result = deduplicator.check_duplicate(fingerprint, "proj-1", mock_session)
            assert result is not None
            # 返回第一个通过阈值的匹配（vid-1 距离=1 < 10）
            assert result["duplicate_of"] == "vid-1"
        finally:
            self._restore_repo(mod, orig)

    def test_no_phashes_skips_video(self, deduplicator, mock_session):
        """已有视频无 phashes 时应被跳过。"""
        existing = self._make_existing_video("vid-1", "md5_a", phashes=[])
        mock_repo = MagicMock()
        mock_repo.list_by_project.return_value = [existing]

        fingerprint = VideoFingerprint(
            md5="md5_b",
            keyframe_phashes=["abcdef01"],
            color_histograms=[],
            duration=10.0,
            resolution=(1280, 720),
        )

        orig, mod = self._patch_repo(mock_repo)
        try:
            result = deduplicator.check_duplicate(fingerprint, "proj-1", mock_session)
            assert result is None
        finally:
            self._restore_repo(mod, orig)

    def test_phash_similarity_formula(self, deduplicator, mock_session):
        """验证相似度公式：similarity = 1.0 - (avg_distance / 64)。"""
        # 使用已知距离的 phash 对
        # "0000000000000000" vs "0000000000000001" → XOR = 1 → 1 bit → distance = 1
        existing = self._make_existing_video("vid-1", "md5_a", phashes=["0000000000000000"])
        mock_repo = MagicMock()
        mock_repo.list_by_project.return_value = [existing]

        fingerprint = VideoFingerprint(
            md5="md5_b",
            keyframe_phashes=["0000000000000001"],  # 1 bit different
            color_histograms=[],
            duration=10.0,
            resolution=(1280, 720),
        )

        orig, mod = self._patch_repo(mock_repo)
        try:
            result = deduplicator.check_duplicate(fingerprint, "proj-1", mock_session)
            assert result is not None
            assert result["duplicate"] is True
            # similarity = 1.0 - (1 / 64) = 0.984375
            assert abs(result["similarity"] - (1.0 - 1.0 / 64)) < 1e-6
        finally:
            self._restore_repo(mod, orig)

    def test_multiple_phashes_avg_distance(self, deduplicator, mock_session):
        """多帧 phash 使用平均最小距离。"""
        # 已有视频有 2 帧 phash
        existing = self._make_existing_video(
            "vid-1", "md5_a",
            phashes=["0000000000000000", "ffffffffffffffff"],
        )
        mock_repo = MagicMock()
        mock_repo.list_by_project.return_value = [existing]

        # 新视频有 1 帧 phash，与第一帧距离=0，与第二帧距离=64
        # min_distance = 0, avg = 0 → 匹配
        fingerprint = VideoFingerprint(
            md5="md5_b",
            keyframe_phashes=["0000000000000000"],
            color_histograms=[],
            duration=10.0,
            resolution=(1280, 720),
        )

        orig, mod = self._patch_repo(mock_repo)
        try:
            result = deduplicator.check_duplicate(fingerprint, "proj-1", mock_session)
            assert result is not None
            assert result["duplicate"] is True
            assert result["similarity"] == 1.0  # avg_distance = 0
        finally:
            self._restore_repo(mod, orig)