111 lines
3.6 KiB
Python
111 lines
3.6 KiB
Python
import base64
|
|
import unittest
|
|
from types import SimpleNamespace
|
|
from unittest.mock import Mock, patch
|
|
|
|
from app.ports.image_gen import ImageResult, TaskStatus
|
|
from app.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("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_successful_generation_persists_completed_status(
|
|
self,
|
|
_acquire_lock_mock,
|
|
_release_lock_mock,
|
|
prompt_service_cls,
|
|
get_image_generator_mock,
|
|
storage_cls,
|
|
get_sync_db_mock,
|
|
):
|
|
chapter = _chapter_stub()
|
|
db = Mock()
|
|
db.execute.return_value.unique.return_value.scalar_one_or_none.return_value = (
|
|
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",
|
|
}
|
|
mock_gen = Mock()
|
|
mock_gen.generate.return_value = ImageResult(
|
|
status=TaskStatus.COMPLETED,
|
|
task_id="",
|
|
image_url="https://provider.example.com/1.png",
|
|
)
|
|
mock_gen.download_image.return_value = _ONE_BY_ONE_PNG
|
|
get_image_generator_mock.return_value = mock_gen
|
|
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")
|