life-echo/api/tests/test_generate_chapter_images_task.py

import unittest
from io import BytesIO
from types import SimpleNamespace
from unittest.mock import Mock, patch

from PIL import Image

from api.tasks import memoir_tasks
from api.tasks.memoir_tasks import generate_chapter_images


def _section_image_record(img_dict):
    """把图片 dict 转成 image_record 用的 SimpleNamespace（可被任务更新属性）。"""
    d = dict(img_dict or {})
    return SimpleNamespace(
        order_index=d.get("index", 0),
        placeholder=d.get("placeholder"),
        description=d.get("description"),
        status=d.get("status"),
        prompt=d.get("prompt"),
        url=d.get("url"),
        storage_key=d.get("storage_key"),
        provider=d.get("provider"),
        style=d.get("style"),
        size=d.get("size"),
        error=d.get("error"),
        retryable=d.get("retryable"),
        created_at=d.get("created_at"),
        updated_at=d.get("updated_at"),
    )


def _chapter_with_sections(sections_data):
    """构造带 sections 的 chapter stub，供 generate_chapter_images 使用（任务从 section.image_record 读/写）。"""
    sections = []
    for i, d in enumerate(sections_data):
        img = d.get("image")
        if img:
            rec = _section_image_record(img)
            sec = SimpleNamespace(
                content=d.get("content", ""),
                image_id="img-%s-%s" % (i, id(rec)),
                image_record=rec,
                order_index=d.get("order_index", i),
            )
        else:
            sec = SimpleNamespace(
                content=d.get("content", ""),
                image_id=None,
                image_record=None,
                order_index=d.get("order_index", i),
            )
        sections.append(sec)
    return SimpleNamespace(
        id="chapter-1",
        user_id="user-1",
        title="童年的夏天",
        category="childhood",
        cover_image=None,
        images=[],
        sections=sections,
    )


def _bind_db_execute_to_chapter(db_mock, chapter):
    """让 db.execute(select(...)).unique().scalar_one_or_none() 返回 chapter。"""
    db_mock.execute.return_value.unique.return_value.scalar_one_or_none.return_value = chapter


class GenerateChapterImagesTaskTest(unittest.TestCase):
    def setUp(self):
        memoir_tasks._REDIS_CLIENTS.clear()

    @patch("api.tasks.memoir_tasks.redis.from_url")
    @patch("api.tasks.memoir_tasks.SessionLocal")
    @patch("api.tasks.memoir_tasks.TencentCosStorageService")
    @patch("api.tasks.memoir_tasks.LiblibImageProvider")
    @patch("api.tasks.memoir_tasks.MemoirImagePromptService")
    def test_generate_chapter_images_skips_when_lock_is_already_held(
        self,
        prompt_service_cls,
        provider_cls,
        storage_cls,
        session_local_cls,
        redis_from_url,
    ):
        chapter = _chapter_with_sections([
            {"content": "那条路我一直记得。", "image": {"index": 0, "placeholder": "{{{{IMAGE:南方小镇的青石板路}}}}", "description": "南方小镇的青石板路", "status": "pending", "url": None}},
        ])
        db = Mock()
        _bind_db_execute_to_chapter(db, chapter)
        session_local_cls.return_value = db
        redis_from_url.return_value.set.return_value = False

        result = generate_chapter_images.run("chapter-1")

        self.assertEqual(result, {"status": "locked"})
        provider_cls.return_value.submit_generation.assert_not_called()
        storage_cls.from_env.return_value.upload_bytes.assert_not_called()
        db.commit.assert_not_called()

    @patch("api.tasks.memoir_tasks.SessionLocal")
    @patch("api.tasks.memoir_tasks.TencentCosStorageService")
    @patch("api.tasks.memoir_tasks.LiblibImageProvider")
    @patch("api.tasks.memoir_tasks.MemoirImagePromptService")
    @patch("api.tasks.memoir_tasks._release_chapter_image_lock")
    @patch("api.tasks.memoir_tasks._acquire_chapter_image_lock", return_value=True)
    def test_generate_chapter_images_retries_when_any_item_generation_fails(
        self,
        _acquire_lock_mock,
        _release_lock_mock,
        prompt_service_cls,
        provider_cls,
        storage_cls,
        session_local_cls,
    ):
        chapter = _chapter_with_sections([
            {"content": "那条路我一直记得。", "image": {"index": 0, "placeholder": "{{{{IMAGE:南方小镇的青石板路}}}}", "description": "南方小镇的青石板路", "status": "pending", "url": None}},
        ])
        db = Mock()
        _bind_db_execute_to_chapter(db, chapter)
        session_local_cls.return_value = db
        prompt_service_cls.return_value.build_prompt.return_value = {
            "prompt": "A serene southern China town",
            "style": "watercolor",
            "size": "1024x1024",
            "prompt_context": "childhood: 童年的夏天",
        }
        provider_cls.return_value.submit_generation.side_effect = RuntimeError("transient provider error")

        retry_error = RuntimeError("retry requested")
        task_self = SimpleNamespace(request=SimpleNamespace(id="task-1"), retry=Mock(side_effect=retry_error))

        with self.assertRaises(RuntimeError) as ctx:
            generate_chapter_images.run.__func__(task_self, "chapter-1")

        self.assertIs(ctx.exception, retry_error)
        self.assertEqual(chapter.sections[0].image_record.status, "failed")
        self.assertEqual(chapter.sections[0].image_record.error, "transient provider error")
        task_self.retry.assert_called_once()
        storage_cls.from_env.return_value.upload_bytes.assert_not_called()

    @patch("api.tasks.memoir_tasks.SessionLocal")
    @patch("api.tasks.memoir_tasks.TencentCosStorageService")
    @patch("api.tasks.memoir_tasks.LiblibImageProvider")
    @patch("api.tasks.memoir_tasks.MemoirImagePromptService")
    @patch("api.tasks.memoir_tasks._release_chapter_image_lock")
    @patch("api.tasks.memoir_tasks._acquire_chapter_image_lock", return_value=True)
    def test_generate_chapter_images_marks_successful_item_completed(
        self,
        _acquire_lock_mock,
        _release_lock_mock,
        prompt_service_cls,
        provider_cls,
        storage_cls,
        session_local_cls,
    ):
        chapter = _chapter_with_sections([
            {"content": "那条路我一直记得。", "image": {"index": 0, "placeholder": "{{{{IMAGE:南方小镇的青石板路}}}}", "description": "南方小镇的青石板路", "status": "pending", "url": None}},
        ])
        db = Mock()
        _bind_db_execute_to_chapter(db, chapter)
        session_local_cls.return_value = db
        prompt_service_cls.return_value.build_prompt.return_value = {
            "prompt": "A serene southern China town",
            "style": "watercolor",
            "size": "1024x1024",
            "prompt_context": "childhood: 童年的夏天",
        }
        provider_inst = provider_cls.return_value
        provider_inst.submit_generation.return_value = {
            "status": "completed",
            "image_url": "https://provider.example.com/1.png",
        }
        png_buffer = BytesIO()
        Image.new("RGB", (1, 1), color="white").save(png_buffer, format="PNG")
        provider_inst.download_image.return_value = png_buffer.getvalue()
        storage_inst = storage_cls.from_env.return_value
        storage_inst.upload_bytes.return_value = "https://cos.example.com/memoirs/u1/c1/0.png"

        generate_chapter_images.run("chapter-1")

        self.assertEqual(chapter.sections[0].image_record.status, "completed")
        self.assertEqual(chapter.sections[0].image_record.storage_key, "memoirs/user-1/chapter-1/0-7e1f860790.png")
        self.assertEqual(chapter.sections[0].image_record.url, "https://cos.example.com/memoirs/u1/c1/0.png")
        self.assertEqual(chapter.sections[0].image_record.prompt, "A serene southern China town")
        provider_inst.close.assert_called_once()
        db.commit.assert_called()

    @patch("api.tasks.memoir_tasks.SessionLocal")
    @patch("api.tasks.memoir_tasks.TencentCosStorageService")
    @patch("api.tasks.memoir_tasks.LiblibImageProvider")
    @patch("api.tasks.memoir_tasks.MemoirImagePromptService")
    @patch("api.tasks.memoir_tasks.MemoirImageSettings.from_env")
    def test_generate_chapter_images_returns_disabled_when_feature_flag_is_off(
        self,
        settings_from_env,
        prompt_service_cls,
        provider_cls,
        storage_cls,
        session_local_cls,
    ):
        chapter = _chapter_with_sections([
            {"content": "那条路我一直记得。", "image": {"index": 0, "placeholder": "{{{{IMAGE:南方小镇的青石板路}}}}", "description": "南方小镇的青石板路", "status": "pending", "url": None}},
        ])
        settings_from_env.return_value = SimpleNamespace(
            enabled=False,
            max_per_chapter=2,
            provider="liblib",
            default_style="watercolor",
            default_size="1024x1024",
            poll_interval_seconds=3,
            max_attempts=20,
            liblib_template_uuid="tpl-uuid",
        )
        db = Mock()
        _bind_db_execute_to_chapter(db, chapter)
        session_local_cls.return_value = db

        result = generate_chapter_images.run("chapter-1")

        self.assertEqual(result, {"status": "disabled"})
        prompt_service_cls.assert_not_called()
        provider_cls.assert_not_called()
        storage_cls.from_env.return_value.upload_bytes.assert_not_called()
        db.commit.assert_not_called()

    @patch("api.tasks.memoir_tasks.SessionLocal")
    @patch("api.tasks.memoir_tasks.TencentCosStorageService")
    @patch("api.tasks.memoir_tasks.LiblibImageProvider")
    @patch("api.tasks.memoir_tasks.MemoirImagePromptService")
    @patch("api.tasks.memoir_tasks._release_chapter_image_lock")
    @patch("api.tasks.memoir_tasks._acquire_chapter_image_lock", return_value=True)
    def test_generate_chapter_images_converts_non_png_payload_before_upload(
        self,
        _acquire_lock_mock,
        _release_lock_mock,
        prompt_service_cls,
        provider_cls,
        storage_cls,
        session_local_cls,
    ):
        chapter = _chapter_with_sections([
            {"content": "那条路我一直记得。", "image": {"index": 0, "placeholder": "{{{{IMAGE:南方小镇的青石板路}}}}", "description": "南方小镇的青石板路", "status": "pending", "url": None}},
        ])
        image_buffer = BytesIO()
        Image.new("RGB", (2, 1), color="white").save(image_buffer, format="JPEG")
        jpeg_bytes = image_buffer.getvalue()

        db = Mock()
        _bind_db_execute_to_chapter(db, chapter)
        session_local_cls.return_value = db
        prompt_service_cls.return_value.build_prompt.return_value = {
            "prompt": "A serene southern China town",
            "style": "watercolor",
            "size": "1024x1024",
            "prompt_context": "childhood: 童年的夏天",
        }
        provider_inst = provider_cls.return_value
        provider_inst.submit_generation.return_value = {
            "status": "completed",
            "image_url": "https://provider.example.com/1.jpg",
        }
        provider_inst.download_image.return_value = jpeg_bytes
        storage_inst = storage_cls.from_env.return_value
        storage_inst.upload_bytes.return_value = "https://cos.example.com/memoirs/u1/c1/0.png"

        generate_chapter_images.run("chapter-1")

        upload_args = storage_inst.upload_bytes.call_args.args
        self.assertTrue(upload_args[0].startswith(b"\x89PNG\r\n\x1a\n"))
        self.assertEqual(upload_args[2], "image/png")

    @patch("api.tasks.memoir_tasks.SessionLocal")
    @patch("api.tasks.memoir_tasks.TencentCosStorageService")
    @patch("api.tasks.memoir_tasks.LiblibImageProvider")
    @patch("api.tasks.memoir_tasks.MemoirImagePromptService")
    @patch("api.tasks.memoir_tasks._release_chapter_image_lock")
    @patch("api.tasks.memoir_tasks._acquire_chapter_image_lock", return_value=True)
    def test_generate_chapter_images_fails_without_retry_on_permanent_cos_error(
        self,
        _acquire_lock_mock,
        _release_lock_mock,
        prompt_service_cls,
        provider_cls,
        storage_cls,
        session_local_cls,
    ):
        chapter = _chapter_with_sections([
            {"content": "那条路我一直记得。", "image": {"index": 0, "placeholder": "{{{{IMAGE:南方小镇的青石板路}}}}", "description": "南方小镇的青石板路", "status": "pending", "url": None}},
        ])
        image_buffer = BytesIO()
        Image.new("RGB", (1, 1), color="white").save(image_buffer, format="PNG")
        png_bytes = image_buffer.getvalue()

        db = Mock()
        _bind_db_execute_to_chapter(db, chapter)
        session_local_cls.return_value = db
        prompt_service_cls.return_value.build_prompt.return_value = {
            "prompt": "A serene southern China town",
            "style": "watercolor",
            "size": "1024x1024",
            "prompt_context": "childhood: 童年的夏天",
        }
        provider_inst = provider_cls.return_value
        provider_inst.submit_generation.return_value = {
            "status": "completed",
            "image_url": "https://provider.example.com/1.png",
        }
        provider_inst.download_image.return_value = png_bytes
        storage_inst = storage_cls.from_env.return_value
        storage_inst.upload_bytes.side_effect = memoir_tasks.CosUploadError(
            "AccessDenied", retryable=False, request_id="req-403"
        )
        task_self = SimpleNamespace(request=SimpleNamespace(id="task-1"), retry=Mock())
        img_rec = chapter.sections[0].image_record

        result = generate_chapter_images.run.__func__(task_self, "chapter-1")

        self.assertEqual(result, {"status": "success"})
        self.assertIsNone(chapter.sections[0].image_id)
        db.delete.assert_called_with(img_rec)
        task_self.retry.assert_not_called()

    @patch("api.tasks.memoir_tasks.SessionLocal")
    @patch("api.tasks.memoir_tasks.TencentCosStorageService")
    @patch("api.tasks.memoir_tasks.LiblibImageProvider")
    @patch("api.tasks.memoir_tasks.MemoirImagePromptService")
    @patch("api.tasks.memoir_tasks._release_chapter_image_lock")
    @patch("api.tasks.memoir_tasks._acquire_chapter_image_lock", return_value=True)
    def test_generate_chapter_images_skips_completed_items_for_idempotency(
        self,
        _acquire_lock_mock,
        _release_lock_mock,
        prompt_service_cls,
        provider_cls,
        storage_cls,
        session_local_cls,
    ):
        chapter = _chapter_with_sections([
            {"content": "那条路我一直记得。", "image": {"index": 0, "placeholder": "{{{{IMAGE:南方小镇的青石板路}}}}", "description": "南方小镇的青石板路", "status": "completed", "url": "https://cos.example.com/already-there.png"}},
        ])
        db = Mock()
        _bind_db_execute_to_chapter(db, chapter)
        session_local_cls.return_value = db

        generate_chapter_images.run("chapter-1")

        provider_cls.return_value.submit_generation.assert_not_called()
        storage_cls.from_env.return_value.upload_bytes.assert_not_called()