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()