421 lines
16 KiB
Python
421 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 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_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 = 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", (1, 1), color="white").save(image_buffer, format="PNG")
|
|
png_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.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())
|
|
|
|
with self.assertRaises(memoir_tasks.PermanentImageGenerationError) as ctx:
|
|
generate_chapter_images.run.__func__(task_self, "chapter-1")
|
|
|
|
self.assertIn("AccessDenied", str(ctx.exception))
|
|
self.assertEqual(chapter.images[0]["status"], "failed")
|
|
self.assertIn("AccessDenied", chapter.images[0]["error"])
|
|
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 = 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()
|