ComfyUI-Lora-Manager/tests/utils/test_exif_utils.py

import json
from typing import Any, Dict

import piexif  # pyright: ignore[reportMissingTypeStubs]
from PIL import Image, PngImagePlugin

from py.utils.exif_utils import ExifUtils


def test_append_recipe_metadata_includes_checkpoint(monkeypatch, tmp_path):
    captured = {}

    monkeypatch.setattr(
        ExifUtils, "extract_image_metadata", staticmethod(lambda _path: None)
    )

    def fake_update_image_metadata(image_path, metadata):
        captured["path"] = image_path
        captured["metadata"] = metadata
        return image_path

    monkeypatch.setattr(
        ExifUtils, "update_image_metadata", staticmethod(fake_update_image_metadata)
    )

    checkpoint = {
        "type": "checkpoint",
        "modelId": 827184,
        "modelVersionId": 2167369,
        "modelName": "WAI-illustrious-SDXL",
        "modelVersionName": "v15.0",
        "hash": "ABC123",
        "file_name": "WAI-illustrious-SDXL",
        "baseModel": "Illustrious",
    }

    recipe_data = {
        "title": "Semi-realism",
        "base_model": "Illustrious",
        "loras": [],
        "tags": [],
        "checkpoint": checkpoint,
    }

    image_path = tmp_path / "image.webp"
    image_path.write_bytes(b"")

    ExifUtils.append_recipe_metadata(str(image_path), recipe_data)

    assert captured["path"] == str(image_path)
    assert captured["metadata"].startswith("Recipe metadata: ")

    payload = json.loads(captured["metadata"].split("Recipe metadata: ", 1)[1])

    assert payload["checkpoint"] == {
        "type": "checkpoint",
        "modelId": 827184,
        "modelVersionId": 2167369,
        "modelName": "WAI-illustrious-SDXL",
        "modelVersionName": "v15.0",
        "hash": "abc123",
        "file_name": "WAI-illustrious-SDXL",
        "baseModel": "Illustrious",
    }
    assert payload["base_model"] == "Illustrious"


def test_append_recipe_metadata_pixel_preserving_webp(tmp_path):
    """pixel_preserving=True must rewrite only the EXIF chunk of a WebP,
    leaving the pixel chunks byte-identical and the old block replaced."""
    img = Image.new("RGB", (64, 48), (120, 30, 200))
    original_params = (
        "masterpiece, best quality\n"
        "Negative prompt: lowres\n"
        "Steps: 20, Sampler: DPM++ 2M Karras, CFG scale: 7, Seed: 1, "
        "Size: 512x768, Model hash: abc123, Model: foo_v1, Clip skip: 2\n"
        ' Recipe metadata: {"title": "Old", "loras": []}'
    )
    exif = piexif.dump(
        {
            "0th": {},
            "Exif": {
                piexif.ExifIFD.UserComment: b"UNICODE\x00"
                + original_params.encode("utf-16be")
            },
        }
    )
    image_path = tmp_path / "recipe.webp"
    img.save(str(image_path), format="WEBP", exif=exif, quality=85)

    with open(image_path, "rb") as fh:
        before = fh.read()

    new_recipe = {
        "title": "New",
        "base_model": "SDXL",
        "loras": [],
        "gen_params": {"steps": 25},
    }
    ExifUtils.append_recipe_metadata(
        str(image_path), new_recipe, pixel_preserving=True
    )

    with open(image_path, "rb") as fh:
        after = fh.read()

    def chunks(data: bytes) -> Dict[bytes, bytes]:
        pos, result = 12, {}
        while pos + 8 <= len(data):
            fourcc = data[pos : pos + 4]
            size = int.from_bytes(data[pos + 4 : pos + 8], "little")
            result[fourcc] = data[pos + 8 : pos + 8 + size]
            pos += 8 + size + (size % 2)
        return result

    before_chunks = chunks(before)
    after_chunks = chunks(after)
    for fourcc, payload in before_chunks.items():
        if fourcc == b"EXIF":
            assert after_chunks[fourcc] != payload, "EXIF must be replaced"
        else:
            assert after_chunks[fourcc] == payload, f"{fourcc} was re-encoded"

    # The appended block is updated; the original parameters stay in front.
    metadata = ExifUtils.extract_image_metadata(str(image_path))
    assert "Steps: 20" in metadata
    assert 'Recipe metadata: {"title": "New"' in metadata
    assert '"title": "Old"' not in metadata


def test_optimize_image_preserves_workflow_when_converting_png_to_webp(tmp_path):
    image_path = tmp_path / "source.png"
    png_info = PngImagePlugin.PngInfo()
    png_info.add_text("parameters", "prompt text\nSteps: 20")
    png_info.add_text("workflow", json.dumps({"nodes": [{"id": 1}]}))

    Image.new("RGB", (64, 32), color="red").save(image_path, pnginfo=png_info)

    optimized_data, extension = ExifUtils.optimize_image(
        str(image_path),
        target_width=32,
        format="webp",
        quality=85,
        preserve_metadata=True,
    )

    optimized_path = tmp_path / f"optimized{extension}"
    optimized_path.write_bytes(optimized_data)

    exif_dict = piexif.load(str(optimized_path))
    assert exif_dict["0th"] is not None
    assert (
        exif_dict["0th"][piexif.ImageIFD.ImageDescription].decode("utf-8")
        == 'Workflow:{"nodes": [{"id": 1}]}'
    )
    assert exif_dict["Exif"] is not None
    user_comment = exif_dict["Exif"][piexif.ExifIFD.UserComment]
    assert user_comment.startswith(b"UNICODE\0")
    assert user_comment[8:].decode("utf-16be") == "prompt text\nSteps: 20"


def test_update_image_metadata_preserves_webp_workflow(tmp_path):
    image_path = tmp_path / "recipe.webp"
    exif_dict = {
        "0th": {
            piexif.ImageIFD.ImageDescription: 'Workflow:{"nodes":[{"id":1}]}',
        },
        "Exif": {
            piexif.ExifIFD.UserComment: b"UNICODE\0"
            + "prompt text".encode("utf-16be")
        },
    }
    Image.new("RGB", (32, 32), color="blue").save(
        image_path, format="WEBP", exif=piexif.dump(exif_dict), quality=85
    )

    ExifUtils.update_image_metadata(
        str(image_path), 'prompt text\nRecipe metadata: {"title":"recipe"}'
    )

    updated_exif = piexif.load(str(image_path))
    assert updated_exif["0th"] is not None
    assert (
        updated_exif["0th"][piexif.ImageIFD.ImageDescription].decode("utf-8")
        == 'Workflow:{"nodes":[{"id":1}]}'
    )
    assert updated_exif["Exif"] is not None
    updated_comment = updated_exif["Exif"][piexif.ExifIFD.UserComment]
    assert (
        updated_comment[8:].decode("utf-16be")
        == 'prompt text\nRecipe metadata: {"title":"recipe"}'
    )


def test_update_image_metadata_preserves_png_workflow(tmp_path):
    image_path = tmp_path / "recipe.png"
    png_info = PngImagePlugin.PngInfo()
    png_info.add_text("parameters", "prompt text")
    png_info.add_text("workflow", '{"nodes":[{"id":1}]}')
    Image.new("RGB", (32, 32), color="green").save(image_path, pnginfo=png_info)

    ExifUtils.update_image_metadata(
        str(image_path), 'prompt text\nRecipe metadata: {"title":"recipe"}'
    )

    with Image.open(image_path) as img:
        assert img.info["workflow"] == '{"nodes":[{"id":1}]}'
        assert (
            img.info["parameters"]
            == 'prompt text\nRecipe metadata: {"title":"recipe"}'
        )


def test_optimize_image_embeds_supplied_workflow_when_source_has_none(tmp_path):
    """Import paths hand the workflow over as data when the preview source is
    metadata-free (CivitAI's optimized rendition); optimize_image must embed
    it while re-encoding, otherwise the recipe loses has_workflow."""
    image_path = tmp_path / "optimized.webp"
    Image.new("RGB", (64, 32), color="red").save(image_path, format="WEBP", quality=85)

    workflow = {"nodes": [{"id": 1}], "last_node_id": 1}
    optimized_data, extension = ExifUtils.optimize_image(
        str(image_path),
        target_width=32,
        format="webp",
        quality=85,
        preserve_metadata=True,
        workflow=workflow,
    )

    optimized_path = tmp_path / f"embedded{extension}"
    optimized_path.write_bytes(optimized_data)

    metadata = ExifUtils._load_structured_metadata(str(optimized_path))
    assert metadata["workflow"] == json.dumps(workflow)


def test_optimize_image_keeps_source_workflow_over_supplied(tmp_path):
    image_path = tmp_path / "source.png"
    png_info = PngImagePlugin.PngInfo()
    png_info.add_text("workflow", '{"nodes":[{"id":7}]}')
    Image.new("RGB", (64, 32), color="red").save(image_path, pnginfo=png_info)

    optimized_data, extension = ExifUtils.optimize_image(
        str(image_path),
        target_width=32,
        format="webp",
        quality=85,
        preserve_metadata=True,
        workflow={"nodes": [{"id": 1}]},
    )

    optimized_path = tmp_path / f"sourcewins{extension}"
    optimized_path.write_bytes(optimized_data)

    metadata = ExifUtils._load_structured_metadata(str(optimized_path))
    assert metadata["workflow"] == '{"nodes":[{"id":7}]}'


def test_embed_workflow_adds_workflow_to_metadata_free_webp(tmp_path):
    image_path = tmp_path / "preview.webp"
    Image.new("RGB", (32, 32), color="blue").save(image_path, format="WEBP", quality=85)

    workflow = json.dumps({"nodes": [{"id": 1}]})
    returned = ExifUtils.embed_workflow(str(image_path), workflow)

    assert returned == str(image_path)
    metadata = ExifUtils._load_structured_metadata(str(image_path))
    assert metadata["workflow"] == workflow
    with Image.open(image_path) as img:
        assert img.size == (32, 32)


def test_embed_workflow_adds_workflow_to_metadata_free_png(tmp_path):
    image_path = tmp_path / "preview.png"
    Image.new("RGB", (32, 32), color="blue").save(image_path)

    workflow = {"nodes": [{"id": 3}]}
    ExifUtils.embed_workflow(str(image_path), workflow)

    metadata = ExifUtils._load_structured_metadata(str(image_path))
    assert metadata["workflow"] == json.dumps(workflow)


def test_embed_workflow_leaves_existing_workflow_untouched(tmp_path):
    image_path = tmp_path / "preview.png"
    png_info = PngImagePlugin.PngInfo()
    png_info.add_text("workflow", '{"nodes":[{"id":9}]}')
    Image.new("RGB", (32, 32), color="green").save(image_path, pnginfo=png_info)

    ExifUtils.embed_workflow(str(image_path), {"nodes": [{"id": 1}]})

    with Image.open(image_path) as img:
        assert img.info["workflow"] == '{"nodes":[{"id":9}]}'


def test_embed_workflow_ignores_unsupported_payloads_and_containers(tmp_path):
    image_path = tmp_path / "preview.webp"
    Image.new("RGB", (16, 16), color="black").save(image_path, format="WEBP")

    # Nothing to embed / unsupported payload types are no-ops.
    assert ExifUtils.embed_workflow(str(image_path), None) == str(image_path)
    assert ExifUtils.embed_workflow(str(image_path), "") == str(image_path)
    assert ExifUtils.embed_workflow(str(image_path), 123) == str(image_path)
    assert ExifUtils._load_structured_metadata(str(image_path))["workflow"] is None

    video_path = tmp_path / "clip.mp4"
    video_path.write_bytes(b"video")
    assert ExifUtils.embed_workflow(str(video_path), {"nodes": []}) == str(video_path)


# --- ISOBMFF / brotli extraction tests ---

import struct

import brotli  # pyright: ignore[reportMissingTypeStubs]


def _build_jxl_with_brob(payload_json: Dict[str, Any]) -> bytes:
    """Build a minimal JXL container with a brob box containing brotli-compressed JSON."""
    # ISOBMFF box 1: JXL signature box (size=12, type='JXL ', signature)
    box1 = struct.pack(">I", 12) + b"JXL " + bytes([0x0d, 0x0a, 0x87, 0x0a])
    # ISOBMFF box 2: ftyp (size=16, type='ftyp', major='jxl ', minor=0)
    box2 = struct.pack(">I", 16) + b"ftyp" + b"jxl " + struct.pack(">I", 0)
    # ISOBMFF box 3: brob — payload is b'comf' + brotli(json)
    compressed = brotli.compress(json.dumps(payload_json).encode("utf-8"))
    brob_payload = b"comf" + compressed
    box3 = struct.pack(">I", 8 + len(brob_payload)) + b"brob" + brob_payload
    return box1 + box2 + box3


def _build_avif_with_brob(payload_json: Dict[str, Any]) -> bytes:
    """Build a minimal AVIF container with a brob box containing brotli-compressed JSON."""
    compressed = brotli.compress(json.dumps(payload_json).encode("utf-8"))
    brob_payload = b"comf" + compressed
    ftyp_box = struct.pack(">I", 20) + b"ftyp" + b"avif" + struct.pack(">I", 0) + b"avif"
    brob_box = struct.pack(">I", 8 + len(brob_payload)) + b"brob" + brob_payload
    return ftyp_box + brob_box


class TestIsobmffBrotliExtraction:
    """Tests for ISOBMFF brotli metadata extraction in ExifUtils."""

    def test_extract_jxl_brotli_happy_path(self, tmp_path):
        """JXL container with valid brob box extracts prompt and workflow."""
        payload = {"prompt": "a cute cat", "workflow": {"nodes": [{"id": 1}]}}
        data = _build_jxl_with_brob(payload)
        path = tmp_path / "test.jxl"
        path.write_bytes(data)

        result = ExifUtils._load_structured_metadata(str(path))

        assert result["prompt"] == "a cute cat"
        assert result["workflow"] == '{"nodes": [{"id": 1}]}'
        assert result["parameters"] is None
        assert result["comment"] is None

    def test_extract_avif_brotli_happy_path(self, tmp_path):
        """AVIF container with valid brob box extracts prompt and workflow."""
        payload = {"prompt": "landscape", "workflow": {"nodes": []}}
        data = _build_avif_with_brob(payload)
        path = tmp_path / "test.avif"
        path.write_bytes(data)

        result = ExifUtils._load_structured_metadata(str(path))

        assert result["prompt"] == "landscape"
        assert result["workflow"] == '{"nodes": []}'

    def test_extract_no_brob_box_returns_none(self, tmp_path):
        """JXL container without a brob box returns None from _extract_isobmff_brotli."""
        # Only JXL signature + ftyp, no brob
        box1 = struct.pack(">I", 12) + b"JXL " + bytes([0x0d, 0x0a, 0x87, 0x0a])
        box2 = struct.pack(">I", 16) + b"ftyp" + b"jxl " + struct.pack(">I", 0)
        path = tmp_path / "test.jxl"
        path.write_bytes(box1 + box2)

        # The low-level extraction should return None (no brob box)
        result = ExifUtils._extract_isobmff_brotli(str(path))
        assert result is None

    def test_extract_corrupt_brob_returns_none(self, tmp_path):
        """Broken brob box payload gracefully returns None."""
        box1 = struct.pack(">I", 12) + b"JXL " + bytes([0x0d, 0x0a, 0x87, 0x0a])
        box2 = struct.pack(">I", 16) + b"ftyp" + b"jxl " + struct.pack(">I", 0)
        # brob with garbage payload that doesn't start with b'comf'
        garbage = b"\xff\xff\xff\xff" * 32
        box3 = struct.pack(">I", 8 + len(garbage)) + b"brob" + garbage
        path = tmp_path / "test.jxl"
        path.write_bytes(box1 + box2 + box3)

        result = ExifUtils._extract_isobmff_brotli(str(path))
        assert result is None

    def test_extract_non_isobmff_file_falls_through(self, tmp_path):
        """A regular PNG file is not processed as ISOBMFF and returns PIL metadata."""
        png_info = PngImagePlugin.PngInfo()
        png_info.add_text("prompt", "from png")
        path = tmp_path / "test.png"
        Image.new("RGB", (4, 4), color="red").save(path, pnginfo=png_info)

        result = ExifUtils._load_structured_metadata(str(path))
        assert result["prompt"] == "from png"

    def test_extract_skip_on_update_and_optimize(self, tmp_path):
        """AVIF/JXL files are skipped for write operations (update/append/optimize)."""
        path = tmp_path / "test.avif"
        path.write_bytes(b"fake avif data")

        # update_image_metadata should return the path unchanged
        result = ExifUtils.update_image_metadata(str(path), "some metadata")
        assert result == str(path)

        # append_recipe_metadata should also skip
        result = ExifUtils.append_recipe_metadata(str(path), {"title": "test"})
        assert result == str(path)

        # optimize_image should passthrough for AVIF/JXL paths
        result_data, ext = ExifUtils.optimize_image(str(path))
        assert ext == ".avif"
        assert result_data == b"fake avif data"

    def test_extract_prompt_as_dict(self, tmp_path):
        """prompt field as dict is JSON-serialized."""
        payload = {"prompt": {"text": "hello", "negative": "bad"}}
        data = _build_jxl_with_brob(payload)
        path = tmp_path / "test.jxl"
        path.write_bytes(data)

        result = ExifUtils._load_structured_metadata(str(path))
        assert result["prompt"] is not None
        assert json.loads(result["prompt"]) == {"text": "hello", "negative": "bad"}

    def test_extract_workflow_as_list(self, tmp_path):
        """workflow field as list is JSON-serialized."""
        payload = {"workflow": [{"id": 1}, {"id": 2}]}
        data = _build_avif_with_brob(payload)
        path = tmp_path / "test.avif"
        path.write_bytes(data)

        result = ExifUtils._load_structured_metadata(str(path))
        assert result["workflow"] is not None
        assert json.loads(result["workflow"]) == [{"id": 1}, {"id": 2}]

    def test_over_decompressed_size_limit(self, tmp_path, monkeypatch):
        """Decompressed data exceeding _BROTLI_MAX_DECOMPRESSED is rejected."""
        # Monkey-patch the limit to a small value to avoid large test data
        monkeypatch.setattr(ExifUtils, "_BROTLI_MAX_DECOMPRESSED", 100)

        large_content = "x" * 200
        payload = {"prompt": large_content}
        data = _build_jxl_with_brob(payload)
        path = tmp_path / "test.jxl"
        path.write_bytes(data)

        # Direct extraction should return None because decompressed size exceeds limit
        result = ExifUtils._extract_isobmff_brotli(str(path))
        assert result is None


# --- get_image_dimensions tests ---


def test_get_image_dimensions_returns_actual_size(tmp_path):
    """(a) A valid image returns its real (width, height)."""
    image_path = tmp_path / "preview.png"
    Image.new("RGB", (64, 32), color="red").save(image_path)

    assert ExifUtils.get_image_dimensions(str(image_path)) == (64, 32)


def test_get_image_dimensions_missing_path_returns_none(tmp_path):
    """(b) A nonexistent path returns None without raising."""
    assert ExifUtils.get_image_dimensions(str(tmp_path / "missing.png")) is None


def test_get_image_dimensions_skips_video_extension_without_pil(tmp_path, monkeypatch):
    """(c) A .mp4 path returns None and never invokes PIL."""
    video_path = tmp_path / "preview.mp4"
    video_path.write_bytes(b"not really a video")

    def fail_if_called(*args, **kwargs):
        raise AssertionError("PIL Image.open must not be called for video paths")

    monkeypatch.setattr("py.utils.exif_utils.Image.open", fail_if_called)

    assert ExifUtils.get_image_dimensions(str(video_path)) is None


def test_get_image_dimensions_corrupt_file_returns_none(tmp_path):
    """(d) A corrupt file returns None without raising."""
    image_path = tmp_path / "corrupt.png"
    image_path.write_bytes(b"\x00\x01\x02\x03 not a real image")

    assert ExifUtils.get_image_dimensions(str(image_path)) is None


def test_get_image_dimensions_cache_key_includes_mtime(tmp_path):
    """(e) Replacing a path with a different-size image returns the new size."""
    image_path = tmp_path / "replaced.png"
    Image.new("RGB", (64, 32), color="red").save(image_path)
    assert ExifUtils.get_image_dimensions(str(image_path)) == (64, 32)

    Image.new("RGB", (100, 50), color="blue").save(image_path)
    assert ExifUtils.get_image_dimensions(str(image_path)) == (100, 50)


def test_get_image_dimensions_skips_unreadable_formats(tmp_path):
    """(f) .avif/.jxl paths return None without raising."""
    for ext in (".avif", ".jxl"):
        image_path = tmp_path / f"preview{ext}"
        image_path.write_bytes(b"fake container data")
        assert ExifUtils.get_image_dimensions(str(image_path)) is None