tps-dashboard/doco/tests/test_video_split.py

# -*- coding: utf-8 -*-
"""
video_split 单元测试
"""

import os
import json
from pathlib import Path

import pytest

# 确保 src 在 path 中
import sys
sys.path.insert(0, str(Path(__file__).parent.parent))

from doco.video_split import (
    hamming_distance,
    format_timestamp,
    build_b_manuscript,
    is_blank_frame,
)


class TestHammingDistance:
    def test_identical_strings(self):
        assert hamming_distance("abc", "abc") == 0

    def test_different_strings(self):
        assert hamming_distance("abc", "abd") == 1

    def test_different_length(self):
        # 长度不同时,返回较长字符串的长度
        assert hamming_distance("abc", "ab") == 3


class TestFormatTimestamp:
    def test_zero(self):
        assert format_timestamp(0) == "0m0s"

    def test_seconds_only(self):
        assert format_timestamp(30000) == "0m30s"  # 30秒

    def test_minutes_and_seconds(self):
        assert format_timestamp(90000) == "1m30s"  # 1分30秒

    def test_longer(self):
        assert format_timestamp(3723000) == "62m3s"  # 62分3秒


class TestBuildBManuscript:
    def test_empty_keyframes(self):
        lines = build_b_manuscript([])
        assert lines == []

    def test_single_frame(self):
        keyframes = [
            {"timestamp_ms": 0, "ocr_text": "测试字幕"}
        ]
        lines = build_b_manuscript(keyframes)
        assert len(lines) == 1
        assert "[0m0s]" in lines[0]
        assert "测试字幕" in lines[0]

    def test_duplicate_text_merged(self):
        """相邻同文本应合并"""
        keyframes = [
            {"timestamp_ms": 0, "ocr_text": "相同"},
            {"timestamp_ms": 1000, "ocr_text": "相同"},
            {"timestamp_ms": 2000, "ocr_text": "不同"},
        ]
        lines = build_b_manuscript(keyframes)
        assert len(lines) == 2

    def test_placeholder_skipped(self):
        """OCR占位文本应跳过"""
        keyframes = [
            {"timestamp_ms": 0, "ocr_text": "[OCR待填充 frame=001.png]"},
            {"timestamp_ms": 1000, "ocr_text": "真实字幕"},
        ]
        lines = build_b_manuscript(keyframes)
        assert len(lines) == 1
        assert "真实字幕" in lines[0]


class TestKeyframesJson:
    """验证 keyframes.json 输出格式"""

    def test_keyframe_structure(self, tmp_path):
        """验证单个关键帧的 JSON 结构"""
        # 模拟关键帧数据
        kf = {
            "frame_index": 1,
            "timestamp_ms": 1000,
            "frame_image_path": str(tmp_path / "frame_0001.png"),
            "phash": "ff00aabb12345678",
            "ocr_text": "测试",
        }

        # 验证字段存在
        assert "frame_index" in kf
        assert "timestamp_ms" in kf
        assert "frame_image_path" in kf
        assert "phash" in kf
        assert "ocr_text" in kf

    def test_keyframes_json_output(self, tmp_path):
        """验证完整 keyframes.json 输出"""
        frames_dir = tmp_path / "frames"
        frames_dir.mkdir()

        # 创建一个测试图片
        test_img = frames_dir / "frame_0001.png"
        test_img.write_bytes(b"fake_png_data")

        keyframes_data = {
            "video_path": str(tmp_path / "video.mp4"),
            "fps_sampled": 1,
            "phash_threshold": 8,
            "keyframes": [
                {
                    "frame_index": 0,
                    "timestamp_ms": 0,
                    "frame_image_path": str(test_img),
                    "phash": "abcd1234",
                    "ocr_text": "首帧",
                }
            ]
        }

        json_path = tmp_path / "keyframes.json"
        with open(json_path, "w", encoding="utf-8") as f:
            json.dump(keyframes_data, f, ensure_ascii=False, indent=2)

        # 验证可读
        with open(json_path, "r", encoding="utf-8") as f:
            loaded = json.load(f)

        assert loaded["fps_sampled"] == 1
        assert loaded["phash_threshold"] == 8
        assert len(loaded["keyframes"]) == 1
        assert loaded["keyframes"][0]["frame_index"] == 0


class TestIsBlankFrame:
    """is_blank_frame 双条件检测单元测试"""

    def test_pure_black_frame(self, tmp_path):
        """全 0 纯黑图:max_brightness=0,white_ratio=0,is_blank=True"""
        from PIL import Image
        img_path = tmp_path / "black.png"
        img = Image.new("L", (100, 100), 0)  # 全黑
        img.save(img_path)

        is_blank, white_ratio, max_brightness = is_blank_frame(img_path)
        assert is_blank is True
        assert max_brightness == 0
        assert white_ratio == 0.0

    def test_subtitle_frame(self, tmp_path):
        """有少量白像素(亮度255,占~1%)的图:is_blank=False"""
        from PIL import Image
        import numpy as np
        img_path = tmp_path / "subtitle.png"
        arr = np.zeros((100, 100), dtype=np.uint8)
        # 1% 白像素(亮度255),其余黑(亮度0)
        arr[:10, :10] = 255  # 10×10=100/10000=1%
        img = Image.fromarray(arr, mode="L")
        img.save(img_path)

        is_blank, white_ratio, max_brightness = is_blank_frame(img_path)
        assert is_blank is False
        assert max_brightness == 255
        assert 0.005 <= white_ratio <= 0.02  # ~1%

    def test_max_brightness_ok_but_ratio_too_low(self, tmp_path):
        """max_brightness>=240 但 white_ratio<0.005 → is_blank=True"""
        from PIL import Image
        import numpy as np
        img_path = tmp_path / "few_pixels.png"
        arr = np.zeros((100, 100), dtype=np.uint8)
        # 只有 0.1% 白像素(不够 0.5%),但 max_brightness=255
        arr[:10, :1] = 255
        img = Image.fromarray(arr, mode="L")
        img.save(img_path)

        is_blank, white_ratio, max_brightness = is_blank_frame(img_path)
        assert is_blank is True
        assert max_brightness == 255
        assert white_ratio < 0.005

    def test_ratio_ok_but_max_brightness_too_low(self, tmp_path):
        """max_brightness<240 → 白像素计数全为0 → is_blank=True"""
        from PIL import Image
        import numpy as np
        img_path = tmp_path / "dim_pixels.png"
        arr = np.zeros((100, 100), dtype=np.uint8)
        # 1% 像素亮度=220(不满足阈值240),其余黑
        # is_blank_frame 用 arr > 240 统计白像素,220 全不满足 → white_ratio=0
        arr[:10, :10] = 220
        img = Image.fromarray(arr, mode="L")
        img.save(img_path)

        is_blank, white_ratio, max_brightness = is_blank_frame(img_path)
        assert is_blank is True
        assert max_brightness == 220
        assert white_ratio == 0.0