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 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 = type( "ChapterStub", (), { "id": "chapter-1", "user_id": "user-1", "title": "童年的夏天", "category": "childhood", "content": "那条路我一直记得。", "images": [ { "index": 0, "placeholder": "{{{{IMAGE:南方小镇的青石板路}}}}", "description": "南方小镇的青石板路", "status": "pending", "url": None, } ], }, )() db = Mock() db.get.return_value = 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 = type( "ChapterStub", (), { "id": "chapter-1", "user_id": "user-1", "title": "童年的夏天", "category": "childhood", "content": "那条路我一直记得。", "images": [ { "index": 0, "placeholder": "{{{{IMAGE:南方小镇的青石板路}}}}", "description": "南方小镇的青石板路", "status": "pending", "url": None, } ], }, )() db = Mock() db.get.return_value = 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.images[0]["status"], "failed") self.assertEqual(chapter.images[0]["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 = type( "ChapterStub", (), { "id": "chapter-1", "user_id": "user-1", "title": "童年的夏天", "category": "childhood", "content": "那条路我一直记得。\n\n{{{{IMAGE:南方小镇的青石板路}}}}", "images": [ { "index": 0, "placeholder": "{{{{IMAGE:南方小镇的青石板路}}}}", "description": "南方小镇的青石板路", "status": "pending", "url": None, } ], }, )() db = Mock() db.get.return_value = 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.images[0]["status"], "completed") self.assertEqual(chapter.images[0]["storage_key"], "memoirs/user-1/chapter-1/0-7e1f860790.png") self.assertEqual(chapter.images[0]["url"], "https://cos.example.com/memoirs/u1/c1/0.png") self.assertEqual(chapter.images[0]["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 = type( "ChapterStub", (), { "id": "chapter-1", "user_id": "user-1", "title": "童年的夏天", "category": "childhood", "content": "那条路我一直记得。", "images": [ { "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() db.get.return_value = chapter session_local_cls.return_value = db result = generate_chapter_images.run("chapter-1") self.assertEqual(result, {"status": "disabled"}) self.assertEqual(chapter.images, []) 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_called_once() @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 = type( "ChapterStub", (), { "id": "chapter-1", "user_id": "user-1", "title": "童年的夏天", "category": "childhood", "content": "那条路我一直记得。\n\n{{{{IMAGE:南方小镇的青石板路}}}}", "images": [ { "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() db.get.return_value = 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_skips_completed_items_for_idempotency( self, _acquire_lock_mock, _release_lock_mock, prompt_service_cls, provider_cls, storage_cls, session_local_cls, ): chapter = type( "ChapterStub", (), { "id": "chapter-1", "user_id": "user-1", "title": "童年的夏天", "category": "childhood", "content": "那条路我一直记得。", "images": [ { "index": 0, "placeholder": "{{{{IMAGE:南方小镇的青石板路}}}}", "description": "南方小镇的青石板路", "status": "completed", "url": "https://cos.example.com/already-there.png", } ], }, )() db = Mock() db.get.return_value = 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()