357 lines
16 KiB
Python
357 lines
16 KiB
Python
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 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 _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("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.MemoirImagePromptService")
|
||
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_sections([
|
||
{"content": "那条路我一直记得。", "image": {"index": 0, "placeholder": "{{{{IMAGE:南方小镇的青石板路}}}}", "description": "南方小镇的青石板路", "status": "pending", "url": None}},
|
||
])
|
||
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.MemoirImagePromptService")
|
||
@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_any_item_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_sections([
|
||
{"content": "那条路我一直记得。", "image": {"index": 0, "placeholder": "{{{{IMAGE:南方小镇的青石板路}}}}", "description": "南方小镇的青石板路", "status": "pending", "url": None}},
|
||
])
|
||
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_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(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("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.MemoirImagePromptService")
|
||
@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_item_completed(
|
||
self,
|
||
_acquire_lock_mock,
|
||
_release_lock_mock,
|
||
prompt_service_cls,
|
||
get_image_generator_mock,
|
||
storage_cls,
|
||
get_sync_db_mock,
|
||
):
|
||
chapter = _chapter_with_sections([
|
||
{"content": "那条路我一直记得。", "image": {"index": 0, "placeholder": "{{{{IMAGE:南方小镇的青石板路}}}}", "description": "南方小镇的青石板路", "status": "pending", "url": None}},
|
||
])
|
||
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_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.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")
|
||
get_image_generator_mock.return_value.generate.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.MemoirImagePromptService")
|
||
@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_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)
|
||
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.MemoirImagePromptService")
|
||
@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_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)
|
||
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_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(
|
||
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/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("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.MemoirImagePromptService")
|
||
@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_sections([
|
||
{"content": "那条路我一直记得。", "image": {"index": 0, "placeholder": "{{{{IMAGE:南方小镇的青石板路}}}}", "description": "南方小镇的青石板路", "status": "pending", "url": None}},
|
||
])
|
||
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_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())
|
||
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("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.MemoirImagePromptService")
|
||
@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_items_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_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)
|
||
get_sync_db_mock.return_value.__enter__.return_value = db
|
||
get_sync_db_mock.return_value.__exit__.return_value = False
|
||
|
||
generate_chapter_images.run("chapter-1")
|
||
|
||
get_image_generator_mock.return_value.generate.assert_not_called()
|
||
storage_cls.from_env.return_value.upload_bytes.assert_not_called()
|