103 lines
3.3 KiB
Python
103 lines
3.3 KiB
Python
|
import base64
|
||
|
|
import unittest
|
||
|
|
from types import SimpleNamespace
|
||
|
|
from unittest.mock import Mock, patch
|
||
|
|
|
||
|
|
from tasks.memoir_tasks import generate_chapter_images
|
||
|
|
|
||
|
|
|
||
|
|
_ONE_BY_ONE_PNG = base64.b64decode(
|
||
|
|
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+aF9sAAAAASUVORK5CYII="
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _image_record(img_dict):
|
||
|
|
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_stub():
|
||
|
|
rec = _image_record(
|
||
|
|
{
|
||
|
|
"index": 0,
|
||
|
|
"placeholder": "{{{{IMAGE:南方小镇的青石板路}}}}",
|
||
|
|
"description": "南方小镇的青石板路",
|
||
|
|
"status": "pending",
|
||
|
|
"url": None,
|
||
|
|
}
|
||
|
|
)
|
||
|
|
section = SimpleNamespace(
|
||
|
|
content="那条路我一直记得。",
|
||
|
|
image_id="image-1",
|
||
|
|
image_record=rec,
|
||
|
|
order_index=0,
|
||
|
|
)
|
||
|
|
return SimpleNamespace(
|
||
|
|
id="chapter-1",
|
||
|
|
user_id="user-1",
|
||
|
|
title="童年的夏天",
|
||
|
|
category="childhood",
|
||
|
|
sections=[section],
|
||
|
|
images=[],
|
||
|
|
cover_image=None,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
class GenerateChapterImagesPersistenceTest(unittest.TestCase):
|
||
|
|
@patch("tasks.memoir_tasks.SessionLocal")
|
||
|
|
@patch("tasks.memoir_tasks.TencentCosStorageService")
|
||
|
|
@patch("tasks.memoir_tasks.LiblibImageProvider")
|
||
|
|
@patch("tasks.memoir_tasks.MemoirImagePromptService")
|
||
|
|
@patch("tasks.memoir_tasks._release_chapter_image_lock")
|
||
|
|
@patch("tasks.memoir_tasks._acquire_chapter_image_lock", return_value=True)
|
||
|
|
def test_successful_generation_persists_completed_status(
|
||
|
|
self,
|
||
|
|
_acquire_lock_mock,
|
||
|
|
_release_lock_mock,
|
||
|
|
prompt_service_cls,
|
||
|
|
provider_cls,
|
||
|
|
storage_cls,
|
||
|
|
session_local_cls,
|
||
|
|
):
|
||
|
|
chapter = _chapter_stub()
|
||
|
|
db = Mock()
|
||
|
|
db.execute.return_value.unique.return_value.scalar_one_or_none.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",
|
||
|
|
}
|
||
|
|
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 = _ONE_BY_ONE_PNG
|
||
|
|
storage_cls.from_env.return_value.upload_bytes.return_value = (
|
||
|
|
"https://cos.example.com/memoirs/user-1/chapter-1/0.png"
|
||
|
|
)
|
||
|
|
|
||
|
|
generate_chapter_images.run("chapter-1")
|
||
|
|
|
||
|
|
record = chapter.sections[0].image_record
|
||
|
|
self.assertEqual(record.status, "completed")
|
||
|
|
self.assertEqual(record.url, "https://cos.example.com/memoirs/user-1/chapter-1/0.png")
|
||
|
|
self.assertEqual(record.prompt, "A serene southern China town")
|