mnemosyne/mnemosyne/library/tests/test_svg_raster.py

"""
Tests for SVG rasterization in the ingest pipeline.

SVG is vector XML: Pillow cannot decode it and the vision stage cannot send it
as a data URI, so it is rendered to PNG at parse time. The cases that matter
are the ones that fail *silently* — Write notes stack their pages inside one
document, and rendering that document whole yields a plausible-looking PNG that
is either blank or an illegible tall strip.
"""

import io
import os
import tempfile

from django.test import TestCase
from PIL import Image

from library.services.parsers import (
    IMAGE_EXTENSIONS,
    DocumentParser,
)
from library.services.svg_raster import (
    SvgRenderError,
    render_svg_pages,
    split_svg_pages,
)

# Geometry copied from a real Write note; the page element is identical across
# both root formats, which is what makes splitting format-agnostic.
PAGE_WIDTH = 1094
PAGE_HEIGHT = 1654
PAGE_PITCH = 1674  # page height + inter-page gap


def write_document(pages: int, root_size: bool = False) -> bytes:
    """
    Build a Write-style document.

    :param pages: Number of stacked pages.
    :param root_size: Emit width/height on the root <svg>. False models
        pre-eeab021 files (the permanent majority), True models newer saves.
    """
    root_attrs = ""
    if root_size:
        total = 10 + pages * PAGE_PITCH
        root_attrs = f' width="{PAGE_WIDTH + 20}" height="{total}"'

    body = "".join(
        f'<svg class="write-page" x="10" y="{10 + i * PAGE_PITCH}" '
        f'width="{PAGE_WIDTH}px" height="{PAGE_HEIGHT}px" '
        f'xmlns="http://www.w3.org/2000/svg">'
        f'<path d="M 100 {100 + i * 40} L {600 + i * 120} {700 + i * 40}" '
        f'stroke="#000000" stroke-width="12" fill="none"/></svg>'
        for i in range(pages)
    )
    return (
        f'<svg id="write-document"{root_attrs} '
        f'xmlns="http://www.w3.org/2000/svg" '
        f'xmlns:xlink="http://www.w3.org/1999/xlink">'
        f'<rect id="write-doc-background" width="100%" height="100%" fill="#808080"/>'
        f'<defs id="write-defs"><style>.write-flat-pen{{fill:none}}</style></defs>'
        f"{body}</svg>"
    ).encode()


GENERIC_SVG = (
    b'<svg xmlns="http://www.w3.org/2000/svg" width="400" height="200" '
    b'viewBox="0 0 400 200">'
    b'<path d="M 20 20 L 380 180" stroke="#000" stroke-width="10"/></svg>'
)


def ink_fraction(png: bytes) -> float:
    """Fraction of non-white pixels — the blank-render canary."""
    grey = Image.open(io.BytesIO(png)).convert("L")
    histogram = grey.histogram()
    return sum(histogram[:200]) / sum(histogram)


class SvgSplitTests(TestCase):
    """Splitting a Write document into per-page SVGs."""

    def test_splits_one_document_per_page(self):
        for pages in (1, 2, 5):
            for root_size in (False, True):
                with self.subTest(pages=pages, root_size=root_size):
                    document = write_document(pages, root_size=root_size)
                    self.assertEqual(len(split_svg_pages(document)), pages)

    def test_both_root_formats_split_identically(self):
        """Root width/height (Write eeab021) must not change the outcome.

        Existing notes are never rewritten, so both formats persist
        indefinitely and have to render the same.
        """
        without = split_svg_pages(write_document(3, root_size=False))
        with_size = split_svg_pages(write_document(3, root_size=True))
        self.assertEqual(len(without), len(with_size))
        self.assertEqual(len(without), 3)

    def test_generic_svg_is_a_single_page(self):
        self.assertEqual(len(split_svg_pages(GENERIC_SVG)), 1)

    def test_unsized_svg_is_rejected_rather_than_guessed(self):
        unsized = b'<svg xmlns="http://www.w3.org/2000/svg"><rect/></svg>'
        with self.assertRaises(SvgRenderError):
            split_svg_pages(unsized)

    def test_script_inside_defs_does_not_desync_the_split(self):
        """Regression: a non-greedy <defs>...</defs> regex mis-parses these."""
        document = write_document(2).replace(
            b'<defs id="write-defs">',
            b'<defs id="write-defs"><script><float value="770" /></script>',
        )
        self.assertEqual(len(split_svg_pages(document)), 2)

    def test_unquoted_root_attributes_are_repaired(self):
        """Write can emit ``width=auto``, which is not valid XML.

        A malformation inside the root tag defeats recovery differently from
        one in the body: libxml2 abandons the whole document and returns a
        bare root, so every page silently disappears.
        """
        document = write_document(2).replace(
            b'<svg id="write-document"',
            b'<svg width=auto height=auto id="write-document"',
        )
        self.assertEqual(len(split_svg_pages(document)), 2)


class SvgRenderTests(TestCase):
    """Rasterizing pages to PNG."""

    def test_renders_one_png_per_page(self):
        pages = render_svg_pages(write_document(3))
        self.assertEqual(len(pages), 3)
        for png, width, height in pages:
            self.assertEqual(Image.open(io.BytesIO(png)).format, "PNG")
            self.assertEqual(max(width, height), 1568)

    def test_rendered_pages_are_not_blank(self):
        """The pre-patch trap: a whole-document render emits ruling, no ink."""
        for root_size in (False, True):
            with self.subTest(root_size=root_size):
                for png, _, _ in render_svg_pages(
                    write_document(3, root_size=root_size)
                ):
                    self.assertGreater(ink_fraction(png), 0.001)

    def test_pages_are_page_shaped_not_a_stacked_strip(self):
        """The post-patch trap: root dimensions span every stacked page.

        Rendering that whole gives one tall strip which, once capped, squashes
        a 5-page note to ~209px wide.
        """
        expected = PAGE_HEIGHT / PAGE_WIDTH
        for root_size in (False, True):
            with self.subTest(root_size=root_size):
                for _, width, height in render_svg_pages(
                    write_document(5, root_size=root_size)
                ):
                    self.assertAlmostEqual(height / width, expected, delta=0.05)

    def test_page_cap_is_respected(self):
        self.assertEqual(len(render_svg_pages(write_document(8), max_pages=5)), 5)


class SvgParserIntegrationTests(TestCase):
    """The parser dispatch — SVG must not reach the Pillow image path."""

    def setUp(self):
        self.parser = DocumentParser()

    def test_svg_is_not_in_image_extensions(self):
        # It was, and PIL.Image.open cannot decode SVG, so ingest always failed.
        self.assertNotIn("svg", IMAGE_EXTENSIONS)

    def test_parse_multipage_svg_yields_one_image_per_page(self):
        with tempfile.NamedTemporaryFile(suffix=".svg", delete=False) as f:
            f.write(write_document(3))
            f.flush()
            path = f.name
        try:
            result = self.parser.parse(path, "svg")
        finally:
            os.unlink(path)

        self.assertEqual(len(result.images), 3)
        self.assertEqual(result.metadata["page_count"], 3)
        self.assertEqual(result.text_blocks, [])
        for index, image in enumerate(result.images):
            # PNG, not svg — the vision stage needs raster for its data URI.
            self.assertEqual(image.ext, "png")
            self.assertEqual(image.source_page, index)
            self.assertGreater(ink_fraction(image.data), 0.001)