ocr_platform/ocr_utils/tests/test_watermark_processor.py

"""WatermarkProcessor 与 preset 单测。"""
import numpy as np
import pytest

from ocr_utils.watermark import WatermarkProcessor, merge_watermark_config, get_preset


def test_merge_cell_preset_has_mask():
    cfg = merge_watermark_config("cell", {"method": "masked_adaptive"})
    assert cfg["method"] == "masked_adaptive"
    assert "mask" in cfg
    assert cfg["mask"]["hough_min_line_length"] < get_preset("page", "masked_adaptive")["mask"][
        "hough_min_line_length"
    ]


def test_processor_disabled_passthrough():
    proc = WatermarkProcessor({"enabled": False}, scope="cell")
    img = np.ones((40, 80, 3), dtype=np.uint8) * 255
    out, stages = proc.process(img)
    assert stages == []
    assert out.shape == img.shape


def test_processor_cell_force_wm():
    proc = WatermarkProcessor(
        {"enabled": True, "method": "masked_adaptive"}, scope="cell"
    )
    img = np.ones((50, 100, 3), dtype=np.uint8) * 240
    cv2 = pytest.importorskip("cv2")
    for x in range(0, 100, 8):
        cv2.line(img, (x, 0), (x + 30, 50), (180, 180, 180), 1)
    out, stages = proc.process(img, force=True)
    assert "wm" in stages
    assert out.shape == img.shape