"""Tests for the CLI entry point.""" from __future__ import annotations import json from typing import TYPE_CHECKING from unittest.mock import MagicMock, patch if TYPE_CHECKING: from pathlib import Path import cv2 import numpy as np import pytest from click.testing import CliRunner from PIL import Image from PIL.PngImagePlugin import PngInfo from remove_ai_watermarks.cli import main @pytest.fixture def runner(): return CliRunner() @pytest.fixture def sample_png(tmp_path: Path) -> Path: """Create a sample PNG for CLI testing.""" img = np.random.randint(0, 255, (200, 200, 3), dtype=np.uint8) path = tmp_path / "input.png" cv2.imwrite(str(path), img) return path def _make_batch_dir(tmp_path: Path, count: int = 3) -> Path: """Create a directory with test images for batch testing.""" input_dir = tmp_path / "input" input_dir.mkdir() for i in range(count): img = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8) cv2.imwrite(str(input_dir / f"img_{i}.png"), img) return input_dir def _make_batch_dir_with_metadata(tmp_path: Path, count: int = 3) -> Path: """Create a directory with PNG images containing AI metadata.""" input_dir = tmp_path / "input" input_dir.mkdir() for i in range(count): img = Image.new("RGB", (64, 64), color=(100 + i, 150, 200)) pnginfo = PngInfo() pnginfo.add_text("parameters", f"Steps: 20, Sampler: Euler, img_{i}") pnginfo.add_text("prompt", "a test landscape") img.save(input_dir / f"img_{i}.png", pnginfo=pnginfo) return input_dir def _mock_invisible_engine(): """Create a mock InvisibleEngine that writes a copy of the input image.""" def _mock_remove_watermark(image_path, output_path=None, **kwargs): out = output_path or image_path.with_stem(image_path.stem + "_clean") out.parent.mkdir(parents=True, exist_ok=True) img = Image.open(image_path) img.save(out) return out mock_engine = MagicMock() mock_engine.remove_watermark.side_effect = _mock_remove_watermark mock_cls = MagicMock(return_value=mock_engine) return mock_cls, mock_engine def _mock_invisible_engine_drops_alpha(): """Mock InvisibleEngine that mimics the real engine's BGR-only output path. The real diffusion-based engine reads with cv2.IMREAD_COLOR and writes a 3-channel result. This mock simulates that so we can regression-test alpha preservation across the ``all`` pipeline. """ def _mock_remove_watermark(image_path, output_path=None, **kwargs): out = output_path or image_path.with_stem(image_path.stem + "_clean") out.parent.mkdir(parents=True, exist_ok=True) bgr = cv2.imread(str(image_path), cv2.IMREAD_COLOR) cv2.imwrite(str(out), bgr) return out mock_engine = MagicMock() mock_engine.remove_watermark.side_effect = _mock_remove_watermark mock_cls = MagicMock(return_value=mock_engine) return mock_cls, mock_engine class TestMainGroup: """Tests for the top-level CLI group.""" def test_help(self, runner): result = runner.invoke(main, ["--help"]) assert result.exit_code == 0 assert "Remove visible and invisible" in result.output def test_version(self, runner): result = runner.invoke(main, ["--version"]) assert result.exit_code == 0 assert "remove-ai-watermarks" in result.output assert "version" in result.output def test_no_command_shows_banner(self, runner): result = runner.invoke(main, []) assert result.exit_code == 0 assert "Remove-AI-Watermarks" in result.output class TestVisibleCommand: """Tests for the 'visible' subcommand.""" def test_visible_help(self, runner): result = runner.invoke(main, ["visible", "--help"]) assert result.exit_code == 0 assert "Gemini watermark" in result.output def test_visible_basic(self, runner, sample_png, tmp_path): output = tmp_path / "clean.png" result = runner.invoke( main, ["visible", str(sample_png), "-o", str(output), "--no-detect"], ) assert result.exit_code == 0 assert output.exists() assert "Saved" in result.output def test_visible_default_output_name(self, runner, sample_png): result = runner.invoke(main, ["visible", str(sample_png), "--no-detect"]) assert result.exit_code == 0 expected = sample_png.with_stem(sample_png.stem + "_clean") assert expected.exists() def test_visible_no_inpaint(self, runner, sample_png, tmp_path): output = tmp_path / "clean.png" result = runner.invoke( main, [ "visible", str(sample_png), "-o", str(output), "--no-inpaint", "--no-detect", ], ) assert result.exit_code == 0 assert output.exists() def test_visible_no_detect(self, runner, sample_png, tmp_path): output = tmp_path / "clean.png" result = runner.invoke( main, ["visible", str(sample_png), "-o", str(output), "--no-detect"], ) assert result.exit_code == 0 def test_visible_nonexistent_file(self, runner): result = runner.invoke(main, ["visible", "/nonexistent/file.png"]) assert result.exit_code != 0 def test_visible_preserves_rgba_transparency(self, runner, tmp_path): """Visible removal on an RGBA PNG must keep the alpha channel, not silently flatten the image onto an opaque background. """ rgba = np.zeros((200, 200, 4), dtype=np.uint8) rgba[:, :, :3] = 200 # light grey foreground rgba[50:150, 50:150, 3] = 255 # opaque square in the middle, rest transparent src = tmp_path / "rgba_in.png" cv2.imwrite(str(src), rgba) output = tmp_path / "rgba_out.png" result = runner.invoke( main, ["visible", str(src), "-o", str(output), "--no-detect"], ) assert result.exit_code == 0, result.output assert output.exists() out = cv2.imread(str(output), cv2.IMREAD_UNCHANGED) assert out.ndim == 3, f"output is not 3D: shape={out.shape}" assert out.shape[2] == 4, f"output is not RGBA: shape={out.shape}" # The transparent corners must remain transparent. assert out[0, 0, 3] == 0 assert out[199, 199, 3] == 0 # The opaque centre remains opaque (the watermark region default is bottom-right, # which doesn't overlap the centre square at 200x200). assert out[100, 100, 3] == 255 def test_visible_clears_alpha_in_watermark_region(self, runner, tmp_path): """When inpainting an RGBA image, the watermark region must be cleared in the alpha channel so the sparkle area becomes transparent, not opaque-black. """ rgba = np.full((200, 200, 4), 255, dtype=np.uint8) # fully opaque white src = tmp_path / "rgba_full.png" cv2.imwrite(str(src), rgba) output = tmp_path / "rgba_cleared.png" result = runner.invoke( main, ["visible", str(src), "-o", str(output), "--no-detect"], ) assert result.exit_code == 0, result.output out = cv2.imread(str(output), cv2.IMREAD_UNCHANGED) assert out.shape[2] == 4 # Default sparkle position is in the bottom-right; alpha there must be 0. from remove_ai_watermarks.gemini_engine import get_watermark_config cfg = get_watermark_config(200, 200) px, py = cfg.get_position(200, 200) size = cfg.logo_size assert out[py + size // 2, px + size // 2, 3] == 0, "alpha in the watermark region was not cleared" def test_visible_rgb_input_stays_rgb(self, runner, sample_png, tmp_path): """Regression: a plain RGB PNG must NOT gain a spurious alpha channel.""" output = tmp_path / "rgb_out.png" result = runner.invoke( main, ["visible", str(sample_png), "-o", str(output), "--no-detect"], ) assert result.exit_code == 0, result.output out = cv2.imread(str(output), cv2.IMREAD_UNCHANGED) assert out.ndim == 3, f"output is not 3D: shape={out.shape}" assert out.shape[2] == 3, f"RGB input produced non-RGB output: shape={out.shape}" class TestInvisibleCommand: """Tests for the 'invisible' subcommand.""" def test_invisible_help(self, runner): result = runner.invoke(main, ["invisible", "--help"]) assert result.exit_code == 0 assert "invisible" in result.output.lower() def test_invisible_basic(self, runner, sample_png, tmp_path): mock_cls, mock_engine = _mock_invisible_engine() output = tmp_path / "clean.png" with ( patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True), patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls), ): result = runner.invoke( main, ["invisible", str(sample_png), "-o", str(output)], ) assert result.exit_code == 0, result.output assert output.exists() mock_engine.remove_watermark.assert_called_once() def test_invisible_default_output(self, runner, sample_png): mock_cls, _mock_engine = _mock_invisible_engine() with ( patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True), patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls), ): result = runner.invoke(main, ["invisible", str(sample_png)]) assert result.exit_code == 0, result.output expected = sample_png.with_stem(sample_png.stem + "_clean") assert expected.exists() def test_invisible_nonexistent_file(self, runner): result = runner.invoke(main, ["invisible", "/nonexistent/file.png"]) assert result.exit_code != 0 class TestAllCommand: """Tests for the 'all' subcommand (full pipeline).""" def test_all_help(self, runner): result = runner.invoke(main, ["all", "--help"]) assert result.exit_code == 0 assert "visible" in result.output.lower() def test_all_basic(self, runner, sample_png, tmp_path): mock_cls, _mock_engine = _mock_invisible_engine() output = tmp_path / "clean.png" with ( patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True), patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls), ): result = runner.invoke( main, ["all", str(sample_png), "-o", str(output)], ) assert result.exit_code == 0, result.output assert output.exists() def test_all_nonexistent_file(self, runner): result = runner.invoke(main, ["all", "/nonexistent/file.png"]) assert result.exit_code != 0 def test_all_preserves_rgba_across_invisible_step(self, runner, tmp_path): """Regression: ``all`` must keep transparency even when the invisible step writes a 3-channel result (as the real diffusion engine does). """ rgba = np.zeros((200, 200, 4), dtype=np.uint8) rgba[:, :, :3] = 200 rgba[50:150, 50:150, 3] = 255 # opaque square; corners transparent src = tmp_path / "rgba_in.png" cv2.imwrite(str(src), rgba) output = tmp_path / "rgba_out.png" mock_cls, _engine = _mock_invisible_engine_drops_alpha() with ( patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True), patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls), patch("remove_ai_watermarks.cli.invisible_available", return_value=True, create=True), patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True), ): result = runner.invoke(main, ["all", str(src), "-o", str(output)]) assert result.exit_code == 0, result.output out = cv2.imread(str(output), cv2.IMREAD_UNCHANGED) assert out.ndim == 3, f"output not 3D: shape={out.shape}" assert out.shape[2] == 4, f"output is not RGBA: shape={out.shape}" assert out[0, 0, 3] == 0 assert out[100, 100, 3] == 255 class TestMetadataCommand: """Tests for the 'metadata' subcommand.""" def test_metadata_help(self, runner): result = runner.invoke(main, ["metadata", "--help"]) assert result.exit_code == 0 def test_metadata_check_clean(self, runner, tmp_clean_png): result = runner.invoke(main, ["metadata", str(tmp_clean_png), "--check"]) assert result.exit_code == 0 assert "No AI metadata" in result.output def test_metadata_check_ai(self, runner, tmp_png_with_ai_metadata): result = runner.invoke(main, ["metadata", str(tmp_png_with_ai_metadata), "--check"]) assert result.exit_code == 0 assert "AI metadata detected" in result.output def test_metadata_remove(self, runner, tmp_png_with_ai_metadata, tmp_path): output = tmp_path / "stripped.png" result = runner.invoke( main, [ "metadata", str(tmp_png_with_ai_metadata), "--remove", "-o", str(output), ], ) assert result.exit_code == 0 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.""" def test_batch_help(self, runner): result = runner.invoke(main, ["batch", "--help"]) assert result.exit_code == 0 def test_batch_empty_dir(self, runner, tmp_path): empty_dir = tmp_path / "empty" empty_dir.mkdir() result = runner.invoke(main, ["batch", str(empty_dir)]) assert result.exit_code == 0 assert "No supported images" in result.output def test_batch_visible_mode(self, runner, tmp_path): input_dir = _make_batch_dir(tmp_path) output_dir = tmp_path / "output" result = runner.invoke( main, ["batch", str(input_dir), "-o", str(output_dir), "--mode", "visible"], ) assert result.exit_code == 0 assert "3 processed" in result.output assert output_dir.exists() assert len(list(output_dir.glob("*.png"))) == 3 def test_batch_metadata_mode(self, runner, tmp_path): input_dir = _make_batch_dir_with_metadata(tmp_path) output_dir = tmp_path / "output" result = runner.invoke( main, ["batch", str(input_dir), "-o", str(output_dir), "--mode", "metadata"], ) assert result.exit_code == 0 assert "3 processed" in result.output assert output_dir.exists() assert len(list(output_dir.glob("*.png"))) == 3 # Verify AI metadata was stripped for out_img in output_dir.glob("*.png"): with Image.open(out_img) as img: assert "parameters" not in img.info def test_batch_invisible_mode(self, runner, tmp_path): input_dir = _make_batch_dir(tmp_path) output_dir = tmp_path / "output" mock_cls, _mock_engine = _mock_invisible_engine() with ( patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True), patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls), patch("remove_ai_watermarks.cli.invisible_available", return_value=True, create=True), patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True), ): result = runner.invoke( main, ["batch", str(input_dir), "-o", str(output_dir), "--mode", "invisible"], ) assert result.exit_code == 0, result.output assert "3 processed" in result.output def test_batch_all_mode(self, runner, tmp_path): input_dir = _make_batch_dir(tmp_path) output_dir = tmp_path / "output" mock_cls, _mock_engine = _mock_invisible_engine() with ( patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True), patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls), patch("remove_ai_watermarks.cli.invisible_available", return_value=True, create=True), patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True), ): result = runner.invoke( main, ["batch", str(input_dir), "-o", str(output_dir), "--mode", "all"], ) assert result.exit_code == 0, result.output assert "3 processed" in result.output def test_batch_all_mode_preserves_rgba(self, runner, tmp_path): """Regression: batch ``all`` must keep transparency across the alpha-dropping invisible step (mirrors test_all_preserves_rgba_...). """ input_dir = tmp_path / "input" input_dir.mkdir() rgba = np.zeros((200, 200, 4), dtype=np.uint8) rgba[:, :, :3] = 200 rgba[50:150, 50:150, 3] = 255 cv2.imwrite(str(input_dir / "rgba.png"), rgba) output_dir = tmp_path / "output" mock_cls, _engine = _mock_invisible_engine_drops_alpha() with ( patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True), patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls), patch("remove_ai_watermarks.cli.invisible_available", return_value=True, create=True), patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True), ): result = runner.invoke( main, ["batch", str(input_dir), "-o", str(output_dir), "--mode", "all"], ) assert result.exit_code == 0, result.output out = cv2.imread(str(output_dir / "rgba.png"), cv2.IMREAD_UNCHANGED) assert out.ndim == 3, f"output not 3D: shape={out.shape}" assert out.shape[2] == 4, f"output is not RGBA: shape={out.shape}" assert out[0, 0, 3] == 0 assert out[100, 100, 3] == 255 def test_batch_default_output_dir(self, runner, tmp_path): input_dir = _make_batch_dir(tmp_path) result = runner.invoke( main, ["batch", str(input_dir), "--mode", "visible"], ) assert result.exit_code == 0 expected_dir = tmp_path / "input_clean" assert expected_dir.exists()