life-echo/api/tests/test_pdf_service_images.py

from io import BytesIO
import unittest
from unittest.mock import AsyncMock, patch, MagicMock

from PIL import Image

from api.services.pdf_service import PDFService
from api.services import pdf_service as pdf_service_module


class PDFServiceImagesTest(unittest.IsolatedAsyncioTestCase):
    @patch("api.services.pdf_service.ReportLabImage")
    @patch("api.services.pdf_service.httpx.AsyncClient")
    @patch("api.services.pdf_service.TencentCosStorageService")
    async def test_generate_pdf_preserves_image_aspect_ratio(
        self,
        storage_cls,
        async_client_cls,
        reportlab_image_cls,
    ):
        image_buffer = BytesIO()
        Image.new("RGB", (2, 1), color="white").save(image_buffer, format="PNG")
        png_bytes = image_buffer.getvalue()
        mock_response = MagicMock()
        mock_response.content = png_bytes
        mock_response.raise_for_status = MagicMock()

        mock_client = AsyncMock()
        mock_client.get.return_value = mock_response
        async_client_cls.return_value.__aenter__ = AsyncMock(return_value=mock_client)
        async_client_cls.return_value.__aexit__ = AsyncMock(return_value=False)
        storage = MagicMock()
        storage.get_download_url.return_value = "https://signed.example.com/0.png?sig=123"
        storage_cls.from_env.return_value = storage
        reportlab_image_cls.return_value = MagicMock()

        service = PDFService()
        book = type("BookStub", (), {"title": "我的回忆录"})()
        chapter = type(
            "ChapterStub",
            (),
            {
                "title": "童年的夏天",
                "content": "{{{{IMAGE:南方小镇的青石板路}}}}",
                "images": [
                    {
                        "index": 0,
                        "placeholder": "{{{{IMAGE:南方小镇的青石板路}}}}",
                        "url": "https://life-echo-dev-1319381411.cos.ap-shanghai.myqcloud.com/memoirs/u1/c1/0.png",
                        "storage_key": "memoirs/u1/c1/0.png",
                        "status": "completed",
                    }
                ],
            },
        )()

        await service.generate_pdf(book, [chapter])

        _, kwargs = reportlab_image_cls.call_args
        self.assertAlmostEqual(kwargs["width"], 5 * 72)
        self.assertAlmostEqual(kwargs["height"], 2.5 * 72)

    @patch("api.services.pdf_service.httpx.AsyncClient")
    @patch("api.services.pdf_service.TencentCosStorageService")
    async def test_generate_pdf_embeds_completed_images_and_removes_placeholders(
        self,
        storage_cls,
        async_client_cls,
    ):
        png_bytes = (
            b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01"
            b"\x08\x02\x00\x00\x00\x90wS\xde\x00\x00\x00\x0cIDATx\x9cc```\x00\x00"
            b"\x00\x04\x00\x01\xf6\x178U\x00\x00\x00\x00IEND\xaeB`\x82"
        )
        mock_response = MagicMock()
        mock_response.content = png_bytes
        mock_response.raise_for_status = MagicMock()

        mock_client = AsyncMock()
        mock_client.get.return_value = mock_response
        async_client_cls.return_value.__aenter__ = AsyncMock(return_value=mock_client)
        async_client_cls.return_value.__aexit__ = AsyncMock(return_value=False)
        storage = MagicMock()
        storage.get_download_url.return_value = "https://signed.example.com/0.png?sig=123"
        storage_cls.from_env.return_value = storage

        service = PDFService()
        book = type("BookStub", (), {"title": "我的回忆录"})()
        chapter = type(
            "ChapterStub",
            (),
            {
                "title": "童年的夏天",
                "content": "那条路我一直记得。\n\n{{{{IMAGE:南方小镇的青石板路}}}}\n\n奶奶常坐在那里。",
                "images": [
                    {
                        "index": 0,
                        "placeholder": "{{{{IMAGE:南方小镇的青石板路}}}}",
                        "url": "https://life-echo-dev-1319381411.cos.ap-shanghai.myqcloud.com/memoirs/u1/c1/0.png",
                        "storage_key": "memoirs/u1/c1/0.png",
                        "status": "completed",
                    }
                ],
            },
        )()

        pdf_bytes = await service.generate_pdf(book, [chapter])

        self.assertGreater(len(pdf_bytes), 100)
        self.assertNotIn(b"IMAGE:", pdf_bytes)
        mock_client.get.assert_called_once_with("https://signed.example.com/0.png?sig=123")

    @patch("api.services.pdf_service.httpx.AsyncClient")
    @patch("api.services.pdf_service.TencentCosStorageService")
    async def test_generate_pdf_skips_private_cos_url_when_signing_fails(
        self,
        storage_cls,
        async_client_cls,
    ):
        mock_client = AsyncMock()
        async_client_cls.return_value.__aenter__ = AsyncMock(return_value=mock_client)
        async_client_cls.return_value.__aexit__ = AsyncMock(return_value=False)
        storage = MagicMock()
        storage.get_download_url.side_effect = pdf_service_module.CosDownloadUrlError(
            "cos unavailable", retryable=True, request_id="req-err"
        )
        storage_cls.from_env.return_value = storage

        service = PDFService()
        book = type("BookStub", (), {"title": "我的回忆录"})()
        chapter = type(
            "ChapterStub",
            (),
            {
                "title": "童年的夏天",
                "content": "那条路我一直记得。\n\n{{{{IMAGE:南方小镇的青石板路}}}}\n\n奶奶常坐在那里。",
                "images": [
                    {
                        "index": 0,
                        "placeholder": "{{{{IMAGE:南方小镇的青石板路}}}}",
                        "url": "https://life-echo-dev-1319381411.cos.ap-shanghai.myqcloud.com/memoirs/u1/c1/0.png",
                        "storage_key": "memoirs/u1/c1/0.png",
                        "status": "completed",
                    }
                ],
            },
        )()

        pdf_bytes = await service.generate_pdf(book, [chapter])

        self.assertGreater(len(pdf_bytes), 100)
        mock_client.get.assert_not_called()