mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-07-10 18:27:04 +02:00
test(identify): cover provenance branches, CLI, sparkle helper
Adds 20 tests around the new provenance path: - identify(): local SD/ComfyUI params -> local-pipeline attribution; visible-sparkle gating at the 0.5 threshold (mocked detector: above, below, unavailable, opt-out); metadata verdict not downgraded by a sparkle hit; OpenAI/SynthID caveats + dedup; ProvenanceReport is JSON-serializable (the CLI --json path); and the honest edge where a C2PA manifest without an AI source marker stays 'unknown'. - CLI 'identify': help, clean PNG, AI PNG platform, valid --json, missing file. - gemini_engine.detect_sparkle_confidence: float in range for a real image, None for an unreadable file. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -2,6 +2,7 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import TYPE_CHECKING
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
@@ -243,6 +244,36 @@ class TestMetadataCommand:
|
||||
assert "stripped" in result.output
|
||||
|
||||
|
||||
class TestIdentifyCommand:
|
||||
"""Tests for the 'identify' subcommand."""
|
||||
|
||||
def test_identify_help(self, runner):
|
||||
result = runner.invoke(main, ["identify", "--help"])
|
||||
assert result.exit_code == 0
|
||||
|
||||
def test_identify_clean_png(self, runner, tmp_clean_png):
|
||||
result = runner.invoke(main, ["identify", str(tmp_clean_png), "--no-visible"])
|
||||
assert result.exit_code == 0
|
||||
assert "unknown" in result.output
|
||||
|
||||
def test_identify_ai_png_reports_platform(self, runner, tmp_png_with_ai_metadata):
|
||||
result = runner.invoke(main, ["identify", str(tmp_png_with_ai_metadata), "--no-visible"])
|
||||
assert result.exit_code == 0
|
||||
assert "AI-generated" in result.output
|
||||
assert "Stable Diffusion" in result.output
|
||||
|
||||
def test_identify_json_is_valid(self, runner, tmp_png_with_ai_metadata):
|
||||
result = runner.invoke(main, ["identify", str(tmp_png_with_ai_metadata), "--no-visible", "--json"])
|
||||
assert result.exit_code == 0
|
||||
payload = json.loads(result.output)
|
||||
assert payload["is_ai_generated"] is True
|
||||
assert payload["confidence"] == "high"
|
||||
|
||||
def test_identify_nonexistent_file(self, runner):
|
||||
result = runner.invoke(main, ["identify", "/nonexistent/file.png"])
|
||||
assert result.exit_code != 0
|
||||
|
||||
|
||||
class TestBatchCommand:
|
||||
"""Tests for the 'batch' subcommand."""
|
||||
|
||||
|
||||
@@ -12,6 +12,7 @@ from remove_ai_watermarks.gemini_engine import (
|
||||
WatermarkPosition,
|
||||
WatermarkSize,
|
||||
_calculate_alpha_map,
|
||||
detect_sparkle_confidence,
|
||||
get_watermark_config,
|
||||
get_watermark_size,
|
||||
)
|
||||
@@ -174,6 +175,21 @@ class TestDetection:
|
||||
assert isinstance(result.gradient_score, float)
|
||||
|
||||
|
||||
class TestDetectSparkleConfidence:
|
||||
"""File-level entry point used by identify.py."""
|
||||
|
||||
def test_returns_float_in_range_for_real_image(self, tmp_image_path):
|
||||
conf = detect_sparkle_confidence(tmp_image_path)
|
||||
assert conf is not None
|
||||
assert 0.0 <= conf <= 1.0
|
||||
|
||||
def test_returns_none_for_unreadable_file(self, tmp_path):
|
||||
# cv2.imread returns None for a non-image; the helper must not raise.
|
||||
bogus = tmp_path / "not_an_image.png"
|
||||
bogus.write_bytes(b"this is not a PNG")
|
||||
assert detect_sparkle_confidence(bogus) is None
|
||||
|
||||
|
||||
# ── Inpainting ──────────────────────────────────────────────────────
|
||||
|
||||
|
||||
|
||||
@@ -6,7 +6,10 @@ against the real committed C2PA / IPTC fixtures in data/samples/.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from dataclasses import asdict
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
@@ -18,6 +21,9 @@ from remove_ai_watermarks.identify import (
|
||||
identify,
|
||||
)
|
||||
|
||||
# Where the lazy import inside identify._visible_sparkle resolves the detector.
|
||||
_SPARKLE_TARGET = "remove_ai_watermarks.gemini_engine.detect_sparkle_confidence"
|
||||
|
||||
SAMPLES_DIR = Path(__file__).resolve().parent.parent / "data" / "samples"
|
||||
|
||||
|
||||
@@ -94,6 +100,16 @@ class TestIdentifyNonPng:
|
||||
r = identify(path, check_visible=False)
|
||||
assert any("SynthID" in w for w in r.watermarks)
|
||||
|
||||
def test_c2pa_without_ai_marker_is_unknown(self, tmp_path: Path):
|
||||
# Adobe signs C2PA on plain Photoshop edits too. Without an AI digital-
|
||||
# source marker, the honest verdict is unknown -- the C2PA watermark is
|
||||
# still listed, but is_ai_generated is not asserted True.
|
||||
path = self._c2pa_jpeg(tmp_path, b"Adobe ... no ai marker here")
|
||||
r = identify(path, check_visible=False)
|
||||
assert r.is_ai_generated is None
|
||||
assert any("C2PA" in w for w in r.watermarks)
|
||||
assert not any("SynthID" in w for w in r.watermarks)
|
||||
|
||||
|
||||
# ── End-to-end verdicts on real fixtures ────────────────────────────
|
||||
|
||||
@@ -134,3 +150,99 @@ class TestIdentifyRealSamples:
|
||||
|
||||
def test_returns_report_dataclass(self):
|
||||
assert isinstance(identify(SAMPLES_DIR / "firefly-1.png", check_visible=False), ProvenanceReport)
|
||||
|
||||
|
||||
# ── Local diffusion parameters (Stable Diffusion / ComfyUI) ─────────
|
||||
|
||||
|
||||
class TestIdentifyLocalParams:
|
||||
"""A PNG carrying SD-style generation params is attributed to a local pipeline."""
|
||||
|
||||
def test_sd_params_attributed_to_local_pipeline(self, tmp_png_with_ai_metadata: Path):
|
||||
r = identify(tmp_png_with_ai_metadata, check_visible=False)
|
||||
assert r.is_ai_generated is True
|
||||
assert r.confidence == "high"
|
||||
assert r.platform is not None
|
||||
assert "Stable Diffusion" in r.platform
|
||||
assert any("generation parameters" in w for w in r.watermarks)
|
||||
|
||||
def test_gen_params_signal_lists_keys(self, tmp_png_with_ai_metadata: Path):
|
||||
r = identify(tmp_png_with_ai_metadata, check_visible=False)
|
||||
signal = next(s for s in r.signals if s.name == "gen_params")
|
||||
assert "parameters" in signal.detail
|
||||
assert signal.confidence == "high"
|
||||
|
||||
def test_clean_png_is_unknown(self, tmp_clean_png: Path):
|
||||
r = identify(tmp_clean_png, check_visible=False)
|
||||
assert r.is_ai_generated is None
|
||||
assert r.platform is None
|
||||
assert r.confidence == "none"
|
||||
assert r.signals == []
|
||||
|
||||
|
||||
# ── Visible-sparkle fallback (mocked detector) ──────────────────────
|
||||
|
||||
|
||||
class TestIdentifyVisibleSparkle:
|
||||
"""The visible-sparkle signal gates on the corpus-tuned threshold (0.5)."""
|
||||
|
||||
def test_above_threshold_promotes_to_medium(self, tmp_clean_png: Path):
|
||||
with patch(_SPARKLE_TARGET, return_value=0.7):
|
||||
r = identify(tmp_clean_png, check_visible=True)
|
||||
assert r.is_ai_generated is True
|
||||
assert r.confidence == "medium"
|
||||
assert r.platform is not None
|
||||
assert "Gemini" in r.platform
|
||||
signal = next(s for s in r.signals if s.name == "visible_sparkle")
|
||||
assert signal.confidence == "medium"
|
||||
|
||||
def test_below_threshold_not_promoted(self, tmp_clean_png: Path):
|
||||
with patch(_SPARKLE_TARGET, return_value=0.4):
|
||||
r = identify(tmp_clean_png, check_visible=True)
|
||||
assert r.is_ai_generated is None
|
||||
assert not any(s.name == "visible_sparkle" for s in r.signals)
|
||||
|
||||
def test_detector_unavailable_does_not_crash(self, tmp_clean_png: Path):
|
||||
with patch(_SPARKLE_TARGET, return_value=None):
|
||||
r = identify(tmp_clean_png, check_visible=True)
|
||||
assert r.is_ai_generated is None
|
||||
assert not any(s.name == "visible_sparkle" for s in r.signals)
|
||||
|
||||
def test_check_visible_false_skips_detector(self, tmp_clean_png: Path):
|
||||
# Even a strong detection is ignored when the caller opts out.
|
||||
with patch(_SPARKLE_TARGET, return_value=0.99) as mock_detect:
|
||||
r = identify(tmp_clean_png, check_visible=False)
|
||||
mock_detect.assert_not_called()
|
||||
assert not any(s.name == "visible_sparkle" for s in r.signals)
|
||||
|
||||
def test_metadata_keeps_high_even_with_sparkle(self, tmp_png_with_ai_metadata: Path):
|
||||
# Metadata verdict (high) is not downgraded by an additional sparkle hit.
|
||||
with patch(_SPARKLE_TARGET, return_value=0.7):
|
||||
r = identify(tmp_png_with_ai_metadata, check_visible=True)
|
||||
assert r.confidence == "high"
|
||||
|
||||
|
||||
# ── Caveats and serialization ───────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.skipif(not SAMPLES_DIR.exists(), reason="data/samples not present")
|
||||
class TestIdentifyCaveats:
|
||||
def test_openai_hedge_caveat_present(self):
|
||||
r = identify(SAMPLES_DIR / "chatgpt-1.png", check_visible=False)
|
||||
assert any("before the rollout" in c for c in r.caveats)
|
||||
|
||||
def test_synthid_proxy_caveat_present(self):
|
||||
r = identify(SAMPLES_DIR / "chatgpt-1.png", check_visible=False)
|
||||
assert any("not locally" in c for c in r.caveats)
|
||||
|
||||
def test_caveats_are_deduplicated(self):
|
||||
r = identify(SAMPLES_DIR / "chatgpt-1.png", check_visible=False)
|
||||
assert len(r.caveats) == len(set(r.caveats))
|
||||
|
||||
|
||||
class TestReportSerializable:
|
||||
def test_report_is_json_serializable(self, tmp_png_with_ai_metadata: Path):
|
||||
# The CLI --json path relies on asdict + json.dumps(default=str).
|
||||
report = identify(tmp_png_with_ai_metadata, check_visible=False)
|
||||
dumped = json.dumps(asdict(report), default=str)
|
||||
assert "is_ai_generated" in dumped
|
||||
|
||||
Reference in New Issue
Block a user