wiltodelta--remove-ai-watermarks/tests/test_api.py

"""High-level convenience API (remove_visible / visible_provenance) and the lazy
top-level re-exports."""

from __future__ import annotations

from pathlib import Path

import numpy as np
import pytest

import remove_ai_watermarks as raiw
from remove_ai_watermarks import api

SAMPLES = Path(__file__).resolve().parents[1] / "data" / "samples"
DOUBAO = SAMPLES / "doubao-1.png"
CHATGPT = SAMPLES / "chatgpt-1.png"


class TestTopLevelExports:
    def test_lazy_reexports_resolve(self):
        assert raiw.remove_visible is api.remove_visible
        assert raiw.visible_provenance is api.visible_provenance

    def test_unknown_attribute_raises(self):
        with pytest.raises(AttributeError):
            _ = raiw.does_not_exist

    def test_bare_import_is_light(self):
        # importing the package must not pull the heavy cv2/torch stack (PEP 562 lazy).
        # Checked in a FRESH interpreter -- another test in this process may already
        # have imported cv2, so an in-process sys.modules check would be flaky.
        import subprocess
        import sys

        code = "import remove_ai_watermarks, sys; print(int(any(m in sys.modules for m in ('cv2','torch'))))"
        out = subprocess.run(  # noqa: S603  -- fixed sys.executable + literal code, no untrusted input
            [sys.executable, "-c", code], check=True, capture_output=True, text=True
        )
        assert out.stdout.strip() == "0", f"bare import pulled a heavy module: {out.stdout!r}"


class TestRemoveVisibleArray:
    def test_array_no_mark_is_noop_copy(self):
        arr = np.zeros((256, 256, 3), np.uint8)
        result, removed = raiw.remove_visible(arr, backend="cv2")
        assert removed == []
        assert result.shape == arr.shape
        assert np.array_equal(result, arr)

    def test_array_accepts_knobs(self):
        arr = np.zeros((256, 256, 3), np.uint8)
        result, removed = raiw.remove_visible(arr, sensitivity="assume_ai", backend="cv2")
        assert removed == []
        assert result.shape == arr.shape

    def test_bad_source_raises(self, tmp_path):
        with pytest.raises(ValueError, match="Could not read image"):
            raiw.remove_visible(tmp_path / "nope.png")


@pytest.mark.skipif(not DOUBAO.exists(), reason="doubao sample not present")
class TestRemoveVisiblePath:
    def test_path_removes_and_writes(self, tmp_path):
        out = tmp_path / "clean.png"
        result, removed = raiw.remove_visible(DOUBAO, out, backend="cv2")
        assert out.exists()
        assert any("Doubao" in lbl for lbl in removed)
        assert result.shape[2] == 3

    def test_path_no_output_returns_without_writing(self, tmp_path):
        # output=None returns the array but writes nothing
        result, _ = raiw.remove_visible(DOUBAO, backend="cv2")
        assert result.ndim == 3


class TestNoOpPreservesOriginal:
    def test_no_mark_copies_original_bytes(self, tmp_path):
        # A clean image (no mark) same-format-out must be copied VERBATIM, not
        # re-encoded -- so a no-op never degrades the original ("work with originals").
        import filecmp

        from PIL import Image

        src = tmp_path / "clean.jpg"
        Image.fromarray(np.full((40, 40, 3), 120, np.uint8), "RGB").save(src, quality=90)
        out = tmp_path / "clean_out.jpg"
        _, removed = raiw.remove_visible(str(src), str(out), sensitivity="strict", backend="cv2")
        assert removed == []
        assert filecmp.cmp(str(src), str(out), shallow=False)  # byte-identical


class TestVisibleProvenance:
    @pytest.mark.skipif(not DOUBAO.exists(), reason="doubao sample not present")
    def test_doubao_tc260_maps_to_bytedance(self):
        prov = raiw.visible_provenance(DOUBAO)
        # TC260 label -> ByteDance family (both doubao and jimeng)
        assert {"doubao", "jimeng"} <= prov

    @pytest.mark.skipif(not CHATGPT.exists(), reason="chatgpt sample not present")
    def test_openai_image_has_no_visible_vendor(self):
        # OpenAI C2PA is not one of the visible-mark vendors -> empty provenance
        assert raiw.visible_provenance(CHATGPT) == frozenset()

    def test_unreadable_path_is_empty(self, tmp_path):
        assert raiw.visible_provenance(tmp_path / "missing.png") == frozenset()


class TestRemoveVisibleOutputPath:
    """Output-path robustness: in-place clean (#3) and a missing output dir (#4)."""

    def _write_clean(self, p: Path) -> None:
        from remove_ai_watermarks import image_io

        image_io.imwrite(str(p), np.full((128, 128, 3), 200, np.uint8))

    def test_inplace_clean_no_crash(self, tmp_path: Path):
        p = tmp_path / "clean.png"
        self._write_clean(p)
        _, removed = raiw.remove_visible(str(p), str(p), backend="cv2")
        assert removed == []
        assert p.exists()

    def test_creates_missing_output_dir(self, tmp_path: Path):
        src = tmp_path / "in.png"
        self._write_clean(src)
        out = tmp_path / "sub" / "out.png"
        raiw.remove_visible(str(src), str(out), backend="cv2")
        assert out.exists()