chatbot-maisonnette/tests/test_attachments_prep.py

"""Tests de la préparation des pièces jointes (images, PDF, fichiers texte). Lancer depuis la racine du projet : python -m unittest"""

import codecs
import io
import random
import unittest
from unittest import mock

from PIL import Image
from pypdf import PdfWriter

from bot import images
from bot.images import MAX_DIMENSION, ImageError, prepare_image
from bot.pdf_files import PdfError, prepare_pdf
from bot.text_files import TextFileError, decode_text_file


def encode(img: Image.Image, fmt: str, **options) -> bytes:
    buffer = io.BytesIO()
    img.save(buffer, format=fmt, **options)
    return buffer.getvalue()


def noise(size: tuple[int, int], mode: str = "RGB") -> Image.Image:
    """Image de bruit (toujours le même) : presque impossible à compresser."""
    channels = len(mode)
    return Image.frombytes(mode, size, random.Random(0).randbytes(size[0] * size[1] * channels))


def opened(data: bytes) -> Image.Image:
    img = Image.open(io.BytesIO(data))
    img.load()
    return img


class PrepareImageTest(unittest.TestCase):
    def test_small_jpeg_and_png_sent_as_is(self):
        for fmt, media_type in (("JPEG", "image/jpeg"), ("PNG", "image/png"),
                                ("GIF", "image/gif"), ("WEBP", "image/webp")):
            with self.subTest(fmt):
                data = encode(Image.new("RGB", (200, 100), "red"), fmt)
                result = prepare_image(data)
                self.assertEqual(result.media_type, media_type)
                self.assertIs(result.data, data)

    def test_exactly_max_dimension_sent_as_is(self):
        data = encode(Image.new("RGB", (MAX_DIMENSION, 100), "red"), "PNG")
        self.assertIs(prepare_image(data).data, data)

    def test_large_png_reduced(self):
        result = prepare_image(encode(Image.new("RGB", (3136, 1000), "blue"), "PNG"))
        self.assertEqual(result.media_type, "image/png")  # sans perte : le texte des captures reste net
        self.assertEqual(opened(result.data).size, (MAX_DIMENSION, 500))

    def test_large_jpeg_reduced(self):
        result = prepare_image(encode(Image.new("RGB", (1000, 3136), "green"), "JPEG"))
        self.assertEqual(result.media_type, "image/jpeg")
        self.assertEqual(opened(result.data).size, (500, MAX_DIMENSION))

    def test_unsupported_formats_converted(self):
        for fmt in ("BMP", "TIFF"):
            with self.subTest(fmt):
                result = prepare_image(encode(Image.new("RGB", (200, 100), "red"), fmt))
                self.assertEqual(result.media_type, "image/png")
                img = opened(result.data)
                self.assertEqual((img.format, img.size), ("PNG", (200, 100)))

    def test_transparency_kept(self):
        result = prepare_image(encode(Image.new("RGBA", (3000, 100), (255, 0, 0, 0)), "PNG"))
        self.assertEqual(result.media_type, "image/png")
        self.assertEqual(opened(result.data).mode, "RGBA")

    def test_heavy_png_becomes_jpeg(self):
        # Limite de l'API abaissée : le PNG la dépasse, le JPEG (avec perte) passe en dessous.
        img = noise((300, 300))
        png, jpeg = encode(img, "PNG", optimize=True), encode(img, "JPEG", quality=45, optimize=True)
        self.assertLess(len(jpeg), len(png))
        with mock.patch.object(images, "API_MAX_BYTES", (len(png) + len(jpeg)) // 2):
            result = prepare_image(encode(img, "PNG"))
        self.assertEqual(result.media_type, "image/jpeg")
        self.assertLessEqual(len(result.data), (len(png) + len(jpeg)) // 2)

    def test_heavy_transparent_png_on_white(self):
        img = noise((300, 300), "RGBA")
        with mock.patch.object(images, "API_MAX_BYTES", len(encode(img, "PNG", optimize=True)) // 2):
            result = prepare_image(encode(img, "PNG"))
        self.assertEqual(result.media_type, "image/jpeg")
        self.assertEqual(opened(result.data).mode, "RGB")

    def test_too_heavy_even_compressed(self):
        with mock.patch.object(images, "API_MAX_BYTES", 1000), self.assertRaises(ImageError) as ctx:
            prepare_image(encode(noise((300, 300)), "PNG"))
        self.assertIn("compresser suffisamment", ctx.exception.user_message)

    def test_not_an_image(self):
        for data in (b"", b"bonjour, ceci n'est pas une image", encode(Image.new("RGB", (50, 50)), "PNG")[:40]):
            with self.subTest(data=data[:10]), self.assertRaises(ImageError) as ctx:
                prepare_image(data)
            self.assertIn("Je n'arrive pas à lire cette image", ctx.exception.user_message)

    def test_too_many_pixels(self):
        # Seuil abaissé pour ne pas fabriquer une vraie image géante.
        with mock.patch.object(Image, "MAX_IMAGE_PIXELS", 1000), self.assertRaises(ImageError) as ctx:
            prepare_image(encode(Image.new("RGB", (100, 100)), "PNG"))
        self.assertIn("beaucoup trop grande", ctx.exception.user_message)


def pdf(pages: int, password: str | None = None, owner_password: str | None = None) -> bytes:
    writer = PdfWriter()
    for _ in range(pages):
        writer.add_blank_page(width=200, height=200)
    if password is not None:
        writer.encrypt(user_password=password, owner_password=owner_password or "propriétaire")
    buffer = io.BytesIO()
    writer.write(buffer)
    return buffer.getvalue()


class PreparePdfTest(unittest.TestCase):
    def test_valid_pdf(self):
        data = pdf(3)
        result = prepare_pdf("cours.pdf", data, max_pages=3)
        self.assertTrue(result.is_pdf)
        self.assertEqual(result.name, "cours.pdf")
        self.assertIs(result.data, data)

    def test_too_many_pages(self):
        with self.assertRaises(PdfError) as ctx:
            prepare_pdf("cours.pdf", pdf(4), max_pages=3)
        self.assertEqual(ctx.exception.user_message,
                         "Ce PDF a 4 pages : je n'en lis que 3 au maximum. Envoie-moi seulement les pages utiles.")

    def test_no_page(self):
        with self.assertRaises(PdfError) as ctx:
            prepare_pdf("vide.pdf", pdf(0), max_pages=3)
        self.assertIn("aucune page", ctx.exception.user_message)

    def test_password(self):
        with self.assertRaises(PdfError) as ctx:
            prepare_pdf("secret.pdf", pdf(1, password="secret"), max_pages=3)
        self.assertIn("mot de passe", ctx.exception.user_message)

    def test_owner_password_only(self):
        # Chiffré sans mot de passe d'ouverture (seulement des droits limités) : lisible.
        self.assertTrue(prepare_pdf("protégé.pdf", pdf(1, password=""), max_pages=3).is_pdf)

    def test_not_a_pdf(self):
        for data in (b"", b"bonjour, ceci n'est pas un PDF", encode(Image.new("RGB", (50, 50)), "PNG")):
            with self.subTest(data=data[:10]), self.assertLogs("chatbot.pdf", "INFO"), \
                    self.assertRaises(PdfError) as ctx:
                prepare_pdf("faux.pdf", data, max_pages=3)
            self.assertEqual(ctx.exception.user_message, "Je n'arrive pas à lire ce PDF.")


class DecodeTextFileTest(unittest.TestCase):
    def test_utf8(self):
        self.assertEqual(decode_text_file("Été à Noël 😀".encode("utf-8")), "Été à Noël 😀")

    def test_utf8_with_bom(self):
        self.assertEqual(decode_text_file(codecs.BOM_UTF8 + "Été".encode("utf-8")), "Été")

    def test_windows_1252(self):
        # « ANSI » du Bloc-notes : pas de l'UTF-8 valide, lu en Windows-1252.
        self.assertEqual(decode_text_file("Été, 5 € œuvre".encode("cp1252")), "Été, 5 € œuvre")

    def test_bytes_undefined_in_windows_1252(self):
        self.assertEqual(decode_text_file(b"caf\xe9 \x81"), "café �")

    def test_utf16_with_bom(self):
        for encoding in ("utf-16-le", "utf-16-be"):
            with self.subTest(encoding):
                bom = codecs.BOM_UTF16_LE if encoding == "utf-16-le" else codecs.BOM_UTF16_BE
                self.assertEqual(decode_text_file(bom + "Été".encode(encoding)), "Été")

    def test_broken_utf16(self):
        with self.assertRaises(TextFileError) as ctx:
            decode_text_file(codecs.BOM_UTF16_LE + b"a\x00b")  # nombre impair d'octets
        self.assertIn("pas du texte lisible", ctx.exception.user_message)

    def test_line_endings_and_spaces(self):
        self.assertEqual(decode_text_file(b"\r\n  ligne 1\r\nligne 2\r\n\r\n"), "ligne 1\nligne 2")

    def test_empty(self):
        for data in (b"", b"  \r\n\t ", codecs.BOM_UTF8):
            with self.subTest(data=data), self.assertRaises(TextFileError) as ctx:
                decode_text_file(data)
            self.assertIn("vide", ctx.exception.user_message)

    def test_binary(self):
        for data in (bytes(range(256)), encode(Image.new("RGB", (50, 50)), "PNG"), "texte".encode("utf-16-le")):
            with self.subTest(data=data[:10]), self.assertRaises(TextFileError) as ctx:
                decode_text_file(data)
            self.assertIn("pas du texte lisible", ctx.exception.user_message)


if __name__ == "__main__":
    unittest.main()