WireViz/tests/test_png_metadata.py

# -*- coding: utf-8 -*-

"""Tests for wv_png_metadata — YAML embedding in PNG iTXT chunks."""

from pathlib import Path

import pytest
from PIL import Image

from wireviz.wv_png_metadata import (
    MAX_YAML_SIZE,
    PNG_KEY_YAML,
    PNG_KEY_VERSION,
    has_yaml_metadata,
    read_yaml_from_png,
    save_yaml_to_png,
)

SAMPLE_YAML = """\
connectors:
  X1:
    pincount: 4
  X2:
    pincount: 4

cables:
  W1:
    wirecount: 4
    length: 1

connections:
  -
    - X1: [1-4]
    - W1: [1-4]
    - X2: [1-4]
"""


def _create_test_png(path: Path, width: int = 10, height: int = 10) -> Path:
    """Create a minimal valid PNG file for testing."""
    im = Image.new("RGB", (width, height), color="white")
    im.save(path)
    return path


# === Round-trip tests ===


class TestRoundTrip:
    def test_basic_round_trip(self, tmp_path):
        """Write YAML to PNG, read it back, verify exact match."""
        png = _create_test_png(tmp_path / "test.png")
        assert save_yaml_to_png(png, SAMPLE_YAML) is True
        result = read_yaml_from_png(png)
        assert result == SAMPLE_YAML

    def test_unicode_round_trip(self, tmp_path):
        """YAML with unicode characters survives embedding."""
        yaml_unicode = "# Kabelbaum\nconnectors:\n  X1:\n    type: 'Stecker \u00f6\u00e4\u00fc\u00df'\n    pincount: 2\n"
        png = _create_test_png(tmp_path / "unicode.png")
        assert save_yaml_to_png(png, yaml_unicode) is True
        result = read_yaml_from_png(png)
        assert result == yaml_unicode
        assert "\u00f6\u00e4\u00fc\u00df" in result

    def test_large_yaml_round_trip(self, tmp_path):
        """YAML at a realistic large size survives embedding."""
        large_yaml = SAMPLE_YAML + ("  # padding line\n" * 10000)
        png = _create_test_png(tmp_path / "large.png")
        assert save_yaml_to_png(png, large_yaml) is True
        result = read_yaml_from_png(png)
        assert result == large_yaml

    def test_multiline_yaml_round_trip(self, tmp_path):
        """YAML with various whitespace patterns survives."""
        yaml_ws = "connectors:\n  X1:\n    notes: |\n      Line 1\n      Line 2\n\n      Line 4 after blank\n    pincount: 2\n"
        png = _create_test_png(tmp_path / "multiline.png")
        assert save_yaml_to_png(png, yaml_ws) is True
        assert read_yaml_from_png(png) == yaml_ws


# === save_yaml_to_png tests ===


class TestSave:
    def test_returns_false_for_missing_file(self, tmp_path):
        """save_yaml_to_png on missing file returns False."""
        missing = tmp_path / "nope.png"
        assert save_yaml_to_png(missing, SAMPLE_YAML) is False

    def test_auto_appends_png_suffix(self, tmp_path):
        """Passing path without .png suffix still works."""
        _create_test_png(tmp_path / "test.png")
        assert save_yaml_to_png(tmp_path / "test", SAMPLE_YAML) is True
        assert has_yaml_metadata(tmp_path / "test.png")

    def test_preserves_existing_metadata(self, tmp_path):
        """Existing PNG text metadata is not destroyed."""
        from PIL.PngImagePlugin import PngInfo

        png_path = tmp_path / "meta.png"
        im = Image.new("RGB", (10, 10), color="white")
        existing = PngInfo()
        existing.add_text("Software", "GraphViz 14.1")
        existing.add_text("Author", "WireViz")
        im.save(png_path, pnginfo=existing)

        save_yaml_to_png(png_path, SAMPLE_YAML)

        with Image.open(png_path) as im:
            im.load()
            assert im.text.get("Software") == "GraphViz 14.1"
            assert im.text.get("Author") == "WireViz"
            assert PNG_KEY_YAML in im.text

    def test_writes_version_tag(self, tmp_path):
        """Metadata includes a version tag for future compatibility."""
        png = _create_test_png(tmp_path / "ver.png")
        save_yaml_to_png(png, SAMPLE_YAML)

        with Image.open(png) as im:
            im.load()
            assert im.text.get(PNG_KEY_VERSION) == "1"

    def test_overwrites_previous_yaml(self, tmp_path):
        """Calling save twice replaces the embedded YAML."""
        png = _create_test_png(tmp_path / "overwrite.png")
        save_yaml_to_png(png, "first: version\n")
        save_yaml_to_png(png, "second: version\n")
        result = read_yaml_from_png(png)
        assert result == "second: version\n"

    def test_atomic_write_no_corruption_on_success(self, tmp_path):
        """File is valid PNG after successful save."""
        png = _create_test_png(tmp_path / "atomic.png")
        save_yaml_to_png(png, SAMPLE_YAML)
        # verify it's still a valid, openable PNG
        with Image.open(png) as im:
            im.load()
            assert im.size == (10, 10)


# === read_yaml_from_png tests ===


class TestRead:
    def test_raises_file_not_found(self, tmp_path):
        """read_yaml_from_png on missing file raises FileNotFoundError."""
        with pytest.raises(FileNotFoundError, match="PNG not found"):
            read_yaml_from_png(tmp_path / "nope.png")

    def test_raises_for_corrupt_file(self, tmp_path):
        """read_yaml_from_png on non-PNG file raises ValueError."""
        corrupt = tmp_path / "corrupt.png"
        corrupt.write_text("this is not a png")
        with pytest.raises(ValueError, match="Cannot read YAML metadata"):
            read_yaml_from_png(corrupt)

    def test_raises_for_png_without_metadata(self, tmp_path):
        """read_yaml_from_png on plain PNG raises ValueError."""
        png = _create_test_png(tmp_path / "plain.png")
        with pytest.raises(ValueError, match="does not contain WireViz YAML"):
            read_yaml_from_png(png)

    def test_raises_for_oversized_yaml(self, tmp_path):
        """read_yaml_from_png rejects YAML exceeding the size limit.

        Pillow's own MAX_TEXT_CHUNK guard fires at the decompression level
        before our application-level size check — defense in depth.
        Both produce ValueError, just with different messages.
        """
        png = _create_test_png(tmp_path / "bomb.png")
        from PIL.PngImagePlugin import PngInfo

        with Image.open(png) as im:
            im.load()
            txt = PngInfo()
            txt.add_itxt(PNG_KEY_YAML, "x" * (MAX_YAML_SIZE + 1), zip=True)
            im.save(png, pnginfo=txt)

        # either our size check or Pillow's decompression limit will fire
        with pytest.raises(ValueError):
            read_yaml_from_png(png)

    def test_auto_appends_png_suffix(self, tmp_path):
        """Passing path without .png suffix still works."""
        png = _create_test_png(tmp_path / "test.png")
        save_yaml_to_png(png, SAMPLE_YAML)
        result = read_yaml_from_png(tmp_path / "test")
        assert result == SAMPLE_YAML


# === has_yaml_metadata tests ===


class TestHasMetadata:
    def test_true_after_embedding(self, tmp_path):
        """has_yaml_metadata returns True after embedding."""
        png = _create_test_png(tmp_path / "test.png")
        save_yaml_to_png(png, SAMPLE_YAML)
        assert has_yaml_metadata(png) is True

    def test_false_for_nonexistent_file(self, tmp_path):
        """has_yaml_metadata returns False for missing file."""
        assert has_yaml_metadata(tmp_path / "nope.png") is False

    def test_false_for_plain_png(self, tmp_path):
        """has_yaml_metadata returns False for PNG without metadata."""
        png = _create_test_png(tmp_path / "plain.png")
        assert has_yaml_metadata(png) is False

    def test_false_for_corrupt_file(self, tmp_path):
        """has_yaml_metadata returns False for non-PNG file."""
        corrupt = tmp_path / "corrupt.png"
        corrupt.write_text("not a png")
        assert has_yaml_metadata(corrupt) is False