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 app.ports.image_gen import ImageResult, TaskStatus
from app.tasks import memoir_tasks
from app.tasks.memoir_tasks import build_cos_key, generate_chapter_images


def _mock_image_generator(
    *,
    image_url: str = "https://provider.example.com/1.png",
    image_bytes: bytes | None = None,
):
    """构造满足 port ImageGenerator 的 mock：generate 返回 ImageResult，download_image 返回 bytes。"""
    if image_bytes is None:
        buf = BytesIO()
        Image.new("RGB", (1, 1), color="white").save(buf, format="PNG")
        image_bytes = buf.getvalue()
    gen = Mock()
    gen.generate.return_value = ImageResult(
        status=TaskStatus.COMPLETED,
        task_id="",
        image_url=image_url,
    )
    gen.download_image.return_value = image_bytes
    return gen


def _chapter_with_cover_memoir_image(
    *,
    cover_status: str = "pending",
    cover_url: str | None = None,
    canonical_markdown: str = "# 童年\n\n那条路我一直记得。",
):
    """stories-first：章节级 MemoirImage（order_index 最小为封面槽位）。"""
    cover_rec = SimpleNamespace(
        id="cover-img-1",
        order_index=0,
        placeholder="",
        description="",
        status=cover_status,
        url=cover_url,
        storage_key=None,
        prompt=None,
        provider=None,
        style=None,
        size=None,
        error=None,
        retryable=None,
        created_at=None,
        updated_at=None,
    )
    return SimpleNamespace(
        id="chapter-1",
        user_id="user-1",
        title="童年的夏天",
        category="childhood",
        canonical_markdown=canonical_markdown,
        cover_image=None,
        images=[cover_rec],
    )


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("app.tasks.memoir_tasks.redis.from_url")
    @patch("app.tasks.memoir_tasks.get_sync_db")
    @patch("app.tasks.memoir_tasks.TencentCosStorageService")
    @patch("app.tasks.memoir_tasks.get_image_generator")
    @patch("app.tasks.memoir_tasks.ImagePromptOrchestrator")
    def test_generate_chapter_images_skips_when_lock_is_already_held(
        self,
        prompt_service_cls,
        get_image_generator_mock,
        storage_cls,
        get_sync_db_mock,
        redis_from_url,
    ):
        chapter = _chapter_with_cover_memoir_image()
        db = Mock()
        _bind_db_execute_to_chapter(db, chapter)
        get_sync_db_mock.return_value.__enter__.return_value = db
        get_sync_db_mock.return_value.__exit__.return_value = False
        redis_from_url.return_value.set.return_value = False

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

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

    @patch("app.tasks.memoir_tasks.get_sync_db")
    @patch("app.tasks.memoir_tasks.TencentCosStorageService")
    @patch("app.tasks.memoir_tasks.get_image_generator")
    @patch("app.tasks.memoir_tasks.ImagePromptOrchestrator")
    @patch("app.tasks.memoir_tasks._release_chapter_image_lock")
    @patch("app.tasks.memoir_tasks._acquire_chapter_image_lock", return_value=True)
    def test_generate_chapter_images_retries_when_cover_generation_fails(
        self,
        _acquire_lock_mock,
        _release_lock_mock,
        prompt_service_cls,
        get_image_generator_mock,
        storage_cls,
        get_sync_db_mock,
    ):
        chapter = _chapter_with_cover_memoir_image()
        cover = chapter.images[0]
        db = Mock()
        _bind_db_execute_to_chapter(db, chapter)
        get_sync_db_mock.return_value.__enter__.return_value = db
        get_sync_db_mock.return_value.__exit__.return_value = False
        prompt_service_cls.return_value.build_cover_prompt.return_value = {
            "prompt": "A serene southern China town",
            "style": "watercolor",
            "size": "1024x1024",
            "prompt_context": "childhood: 童年的夏天",
        }
        get_image_generator_mock.return_value.generate.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(cover.status, "failed")
        self.assertEqual(cover.error, "transient provider error")
        task_self.retry.assert_called_once()
        storage_cls.from_env.return_value.upload_bytes.assert_not_called()

    @patch("app.tasks.memoir_tasks.get_sync_db")
    @patch("app.tasks.memoir_tasks.TencentCosStorageService")
    @patch("app.tasks.memoir_tasks.get_image_generator")
    @patch("app.tasks.memoir_tasks.ImagePromptOrchestrator")
    @patch("app.tasks.memoir_tasks._release_chapter_image_lock")
    @patch("app.tasks.memoir_tasks._acquire_chapter_image_lock", return_value=True)
    def test_generate_chapter_images_marks_successful_cover_completed(
        self,
        _acquire_lock_mock,
        _release_lock_mock,
        prompt_service_cls,
        get_image_generator_mock,
        storage_cls,
        get_sync_db_mock,
    ):
        chapter = _chapter_with_cover_memoir_image()
        cover = chapter.images[0]
        db = Mock()
        _bind_db_execute_to_chapter(db, chapter)
        get_sync_db_mock.return_value.__enter__.return_value = db
        get_sync_db_mock.return_value.__exit__.return_value = False
        prompt_data = {
            "prompt": "A serene southern China town",
            "style": "watercolor",
            "size": "1024x1024",
            "prompt_context": "childhood: 童年的夏天",
        }
        prompt_service_cls.return_value.build_cover_prompt.return_value = prompt_data
        get_image_generator_mock.return_value = _mock_image_generator()
        storage_inst = storage_cls.from_env.return_value
        storage_inst.upload_bytes.return_value = (
            "https://cos.example.com/memoirs/u1/c1/cover.png"
        )

        generate_chapter_images.run("chapter-1")

        self.assertEqual(cover.status, "completed")
        expected_key = build_cos_key(
            "user-1", "chapter-1", "cover", prompt_data["prompt"]
        )
        self.assertEqual(cover.storage_key, expected_key)
        self.assertEqual(
            cover.url,
            "https://cos.example.com/memoirs/u1/c1/cover.png",
        )
        self.assertEqual(cover.prompt, "A serene southern China town")
        get_image_generator_mock.return_value.generate.assert_called_once()
        prompt_service_cls.return_value.build_cover_prompt.assert_called_once()
        db.commit.assert_called()

    @patch("app.tasks.memoir_tasks.get_sync_db")
    @patch("app.tasks.memoir_tasks.TencentCosStorageService")
    @patch("app.tasks.memoir_tasks.get_image_generator")
    @patch("app.tasks.memoir_tasks.ImagePromptOrchestrator")
    @patch("app.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,
        get_image_generator_mock,
        storage_cls,
        get_sync_db_mock,
    ):
        chapter = _chapter_with_cover_memoir_image()
        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)
        get_sync_db_mock.return_value.__enter__.return_value = db
        get_sync_db_mock.return_value.__exit__.return_value = False

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

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

    @patch("app.tasks.memoir_tasks.get_sync_db")
    @patch("app.tasks.memoir_tasks.TencentCosStorageService")
    @patch("app.tasks.memoir_tasks.get_image_generator")
    @patch("app.tasks.memoir_tasks.ImagePromptOrchestrator")
    @patch("app.tasks.memoir_tasks._release_chapter_image_lock")
    @patch("app.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,
        get_image_generator_mock,
        storage_cls,
        get_sync_db_mock,
    ):
        chapter = _chapter_with_cover_memoir_image()
        db = Mock()
        _bind_db_execute_to_chapter(db, chapter)
        get_sync_db_mock.return_value.__enter__.return_value = db
        get_sync_db_mock.return_value.__exit__.return_value = False
        prompt_service_cls.return_value.build_cover_prompt.return_value = {
            "prompt": "A serene southern China town",
            "style": "watercolor",
            "size": "1024x1024",
            "prompt_context": "childhood: 童年的夏天",
        }
        image_buffer = BytesIO()
        Image.new("RGB", (2, 1), color="white").save(image_buffer, format="JPEG")
        jpeg_bytes = image_buffer.getvalue()

        get_image_generator_mock.return_value = _mock_image_generator(
            image_url="https://provider.example.com/1.jpg",
            image_bytes=jpeg_bytes,
        )
        storage_inst = storage_cls.from_env.return_value
        storage_inst.upload_bytes.return_value = (
            "https://cos.example.com/memoirs/u1/c1/cover.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("app.tasks.memoir_tasks.get_sync_db")
    @patch("app.tasks.memoir_tasks.TencentCosStorageService")
    @patch("app.tasks.memoir_tasks.get_image_generator")
    @patch("app.tasks.memoir_tasks.ImagePromptOrchestrator")
    @patch("app.tasks.memoir_tasks._release_chapter_image_lock")
    @patch("app.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,
        get_image_generator_mock,
        storage_cls,
        get_sync_db_mock,
    ):
        chapter = _chapter_with_cover_memoir_image()
        cover = chapter.images[0]
        db = Mock()
        _bind_db_execute_to_chapter(db, chapter)
        get_sync_db_mock.return_value.__enter__.return_value = db
        get_sync_db_mock.return_value.__exit__.return_value = False
        prompt_service_cls.return_value.build_cover_prompt.return_value = {
            "prompt": "A serene southern China town",
            "style": "watercolor",
            "size": "1024x1024",
            "prompt_context": "childhood: 童年的夏天",
        }
        get_image_generator_mock.return_value = _mock_image_generator()
        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())

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

        self.assertEqual(result, {"status": "success"})
        db.delete.assert_called_with(cover)
        task_self.retry.assert_not_called()

    @patch("app.tasks.memoir_tasks.get_sync_db")
    @patch("app.tasks.memoir_tasks.TencentCosStorageService")
    @patch("app.tasks.memoir_tasks.get_image_generator")
    @patch("app.tasks.memoir_tasks.ImagePromptOrchestrator")
    @patch("app.tasks.memoir_tasks._release_chapter_image_lock")
    @patch("app.tasks.memoir_tasks._acquire_chapter_image_lock", return_value=True)
    def test_generate_chapter_images_skips_completed_cover_for_idempotency(
        self,
        _acquire_lock_mock,
        _release_lock_mock,
        prompt_service_cls,
        get_image_generator_mock,
        storage_cls,
        get_sync_db_mock,
    ):
        chapter = _chapter_with_cover_memoir_image(
            cover_status="completed",
            cover_url="https://cos.example.com/already-there.png",
        )
        db = Mock()
        _bind_db_execute_to_chapter(db, chapter)
        get_sync_db_mock.return_value.__enter__.return_value = db
        get_sync_db_mock.return_value.__exit__.return_value = False

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

        self.assertEqual(result, {"status": "no_images"})
        get_image_generator_mock.return_value.generate.assert_not_called()
        storage_cls.from_env.return_value.upload_bytes.assert_not_called()