284 lines
11 KiB
Python
284 lines
11 KiB
Python
import os
|
||
import unittest
|
||
import unittest.mock
|
||
|
||
from api.tasks.memoir_tasks import initialize_chapter_images
|
||
|
||
|
||
class MemoirImageBootstrapTest(unittest.TestCase):
|
||
def test_initialize_chapter_images_keeps_only_completed_assets_when_disabled(self):
|
||
chapter = type(
|
||
"ChapterStub",
|
||
(),
|
||
{
|
||
"id": "chapter-1",
|
||
"title": "童年的夏天",
|
||
"category": "childhood",
|
||
"content": "那条路我一直记得。",
|
||
"images": [
|
||
{
|
||
"index": 0,
|
||
"placeholder": "{{IMAGE:南方小镇的青石板路}}",
|
||
"description": "南方小镇的青石板路",
|
||
"status": "completed",
|
||
"url": "https://cos.example.com/existing.png",
|
||
},
|
||
{
|
||
"index": 1,
|
||
"placeholder": "{{IMAGE:奶奶坐在院子里的藤椅上}}",
|
||
"description": "奶奶坐在院子里的藤椅上",
|
||
"status": "pending",
|
||
"url": None,
|
||
},
|
||
],
|
||
},
|
||
)()
|
||
|
||
with unittest.mock.patch.dict(os.environ, {"MEMOIR_IMAGE_ENABLED": "false"}, clear=False):
|
||
assets = initialize_chapter_images(chapter)
|
||
|
||
self.assertEqual(len(assets), 1)
|
||
self.assertEqual(assets[0]["status"], "completed")
|
||
self.assertEqual(assets[0]["url"], "https://cos.example.com/existing.png")
|
||
self.assertEqual(chapter.images, assets)
|
||
|
||
def test_initialize_chapter_images_sets_pending_assets_when_enabled(self):
|
||
chapter = type(
|
||
"ChapterStub",
|
||
(),
|
||
{
|
||
"id": "chapter-1",
|
||
"title": "童年的夏天",
|
||
"category": "childhood",
|
||
"content": "那条路我一直记得。\n\n{{{{IMAGE:南方小镇的青石板路}}}}",
|
||
"images": [],
|
||
},
|
||
)()
|
||
|
||
with unittest.mock.patch.dict(os.environ, {"MEMOIR_IMAGE_ENABLED": "true"}, clear=False):
|
||
assets = initialize_chapter_images(chapter)
|
||
|
||
self.assertEqual(len(assets), 1)
|
||
self.assertEqual(assets[0]["status"], "pending")
|
||
|
||
def test_initialize_chapter_images_preserves_completed_assets_and_adds_only_new_placeholders(self):
|
||
chapter = type(
|
||
"ChapterStub",
|
||
(),
|
||
{
|
||
"id": "chapter-1",
|
||
"title": "童年的夏天",
|
||
"category": "childhood",
|
||
"content": (
|
||
"那条路我一直记得。\n\n"
|
||
"{{{{IMAGE:南方小镇的青石板路}}}}\n\n"
|
||
"奶奶总坐在门口。\n\n"
|
||
"{{{{IMAGE:奶奶坐在院子里的藤椅上}}}}"
|
||
),
|
||
"images": [
|
||
{
|
||
"index": 0,
|
||
"placeholder": "{{{{IMAGE:南方小镇的青石板路}}}}",
|
||
"description": "南方小镇的青石板路",
|
||
"prompt": "A serene southern China town",
|
||
"url": "https://cos.example.com/existing.png",
|
||
"status": "completed",
|
||
"provider": "liblib",
|
||
"style": "watercolor",
|
||
"size": "1024x1024",
|
||
"error": None,
|
||
"created_at": "2026-03-10T10:00:00Z",
|
||
"updated_at": "2026-03-10T10:00:00Z",
|
||
}
|
||
],
|
||
},
|
||
)()
|
||
|
||
with unittest.mock.patch.dict(os.environ, {"MEMOIR_IMAGE_ENABLED": "true"}, clear=False):
|
||
assets = initialize_chapter_images(chapter)
|
||
|
||
self.assertEqual(len(assets), 2)
|
||
self.assertEqual(assets[0]["status"], "completed")
|
||
self.assertEqual(assets[0]["url"], "https://cos.example.com/existing.png")
|
||
self.assertEqual(assets[1]["status"], "pending")
|
||
self.assertEqual(assets[1]["description"], "奶奶坐在院子里的藤椅上")
|
||
|
||
def test_initialize_chapter_images_accepts_double_brace_placeholders(self):
|
||
chapter = type(
|
||
"ChapterStub",
|
||
(),
|
||
{
|
||
"id": "chapter-1",
|
||
"title": "童年的夏天",
|
||
"category": "childhood",
|
||
"content": "开头。\n\n{{IMAGE:1938年初的上海弄堂口，冬日萧瑟}}\n\n结尾。",
|
||
"images": [],
|
||
},
|
||
)()
|
||
|
||
with unittest.mock.patch.dict(os.environ, {"MEMOIR_IMAGE_ENABLED": "true"}, clear=False):
|
||
assets = initialize_chapter_images(chapter)
|
||
|
||
self.assertEqual(len(assets), 1)
|
||
self.assertEqual(assets[0]["status"], "pending")
|
||
self.assertEqual(assets[0]["placeholder"], "{{IMAGE:1938年初的上海弄堂口，冬日萧瑟}}")
|
||
|
||
def test_initialize_chapter_images_normalizes_invalid_existing_asset_status(self):
|
||
chapter = type(
|
||
"ChapterStub",
|
||
(),
|
||
{
|
||
"id": "chapter-1",
|
||
"title": "童年的夏天",
|
||
"category": "childhood",
|
||
"content": "开头。\n\n{{IMAGE:南方小镇的青石板路}}\n\n结尾。",
|
||
"images": [
|
||
{
|
||
"index": 0,
|
||
"placeholder": "{{IMAGE:南方小镇的青石板路}}",
|
||
"description": "南方小镇的青石板路",
|
||
"status": "mystery",
|
||
}
|
||
],
|
||
},
|
||
)()
|
||
|
||
with unittest.mock.patch.dict(os.environ, {"MEMOIR_IMAGE_ENABLED": "true"}, clear=False):
|
||
assets = initialize_chapter_images(chapter)
|
||
|
||
self.assertEqual(len(assets), 1)
|
||
self.assertEqual(assets[0]["status"], "failed")
|
||
self.assertEqual(assets[0]["error"], "invalid image status: mystery")
|
||
|
||
def test_initialize_chapter_images_preserves_existing_completed_assets_beyond_effective_max(self):
|
||
chapter = type(
|
||
"ChapterStub",
|
||
(),
|
||
{
|
||
"id": "chapter-1",
|
||
"title": "童年的夏天",
|
||
"category": "childhood",
|
||
"content": (
|
||
"{{IMAGE:南方小镇的青石板路}}\n"
|
||
"{{IMAGE:奶奶坐在院子里的藤椅上}}\n"
|
||
"{{IMAGE:门前的老槐树}}"
|
||
),
|
||
"images": [
|
||
{
|
||
"index": 0,
|
||
"placeholder": "{{IMAGE:南方小镇的青石板路}}",
|
||
"description": "南方小镇的青石板路",
|
||
"status": "completed",
|
||
"url": "https://cos.example.com/1.png",
|
||
},
|
||
{
|
||
"index": 1,
|
||
"placeholder": "{{IMAGE:奶奶坐在院子里的藤椅上}}",
|
||
"description": "奶奶坐在院子里的藤椅上",
|
||
"status": "completed",
|
||
"url": "https://cos.example.com/2.png",
|
||
},
|
||
{
|
||
"index": 2,
|
||
"placeholder": "{{IMAGE:门前的老槐树}}",
|
||
"description": "门前的老槐树",
|
||
"status": "completed",
|
||
"url": "https://cos.example.com/3.png",
|
||
},
|
||
],
|
||
},
|
||
)()
|
||
|
||
with unittest.mock.patch.dict(
|
||
os.environ,
|
||
{
|
||
"MEMOIR_IMAGE_ENABLED": "true",
|
||
"MEMOIR_IMAGE_MAX_PER_CHAPTER": "2",
|
||
"MEMOIR_IMAGE_CHARS_PER_EXTRA": "99999",
|
||
"MEMOIR_IMAGE_MAX_CAP": "8",
|
||
},
|
||
clear=False,
|
||
):
|
||
assets = initialize_chapter_images(chapter)
|
||
|
||
self.assertEqual(len(assets), 3)
|
||
self.assertEqual(
|
||
[asset["placeholder"] for asset in assets],
|
||
[
|
||
"{{IMAGE:南方小镇的青石板路}}",
|
||
"{{IMAGE:奶奶坐在院子里的藤椅上}}",
|
||
"{{IMAGE:门前的老槐树}}",
|
||
],
|
||
)
|
||
self.assertTrue(all(asset["status"] == "completed" for asset in assets))
|
||
|
||
def test_initialize_chapter_images_increases_limit_for_long_content(self):
|
||
chapter = type(
|
||
"ChapterStub",
|
||
(),
|
||
{
|
||
"id": "chapter-1",
|
||
"title": "童年的夏天",
|
||
"category": "childhood",
|
||
"content": (
|
||
("很长的正文" * 800)
|
||
+ "\n{{IMAGE:南方小镇的青石板路}}"
|
||
+ "\n{{IMAGE:奶奶坐在院子里的藤椅上}}"
|
||
+ "\n{{IMAGE:门前的老槐树}}"
|
||
+ "\n{{IMAGE:夏夜的晒谷场}}"
|
||
),
|
||
"images": [],
|
||
},
|
||
)()
|
||
|
||
with unittest.mock.patch.dict(
|
||
os.environ,
|
||
{
|
||
"MEMOIR_IMAGE_ENABLED": "true",
|
||
"MEMOIR_IMAGE_MAX_PER_CHAPTER": "2",
|
||
"MEMOIR_IMAGE_CHARS_PER_EXTRA": "1000",
|
||
"MEMOIR_IMAGE_MAX_CAP": "8",
|
||
},
|
||
clear=False,
|
||
):
|
||
assets = initialize_chapter_images(chapter)
|
||
|
||
self.assertEqual(len(assets), 4)
|
||
self.assertTrue(all(asset["status"] == "pending" for asset in assets))
|
||
|
||
def test_initialize_chapter_images_caps_dynamic_limit_at_max_images_cap(self):
|
||
chapter = type(
|
||
"ChapterStub",
|
||
(),
|
||
{
|
||
"id": "chapter-1",
|
||
"title": "童年的夏天",
|
||
"category": "childhood",
|
||
"content": (
|
||
("很长的正文" * 1600)
|
||
+ "\n{{IMAGE:图1}}"
|
||
+ "\n{{IMAGE:图2}}"
|
||
+ "\n{{IMAGE:图3}}"
|
||
+ "\n{{IMAGE:图4}}"
|
||
+ "\n{{IMAGE:图5}}"
|
||
+ "\n{{IMAGE:图6}}"
|
||
),
|
||
"images": [],
|
||
},
|
||
)()
|
||
|
||
with unittest.mock.patch.dict(
|
||
os.environ,
|
||
{
|
||
"MEMOIR_IMAGE_ENABLED": "true",
|
||
"MEMOIR_IMAGE_MAX_PER_CHAPTER": "2",
|
||
"MEMOIR_IMAGE_CHARS_PER_EXTRA": "1000",
|
||
"MEMOIR_IMAGE_MAX_CAP": "4",
|
||
},
|
||
clear=False,
|
||
):
|
||
assets = initialize_chapter_images(chapter)
|
||
|
||
self.assertEqual(len(assets), 4)
|
||
self.assertEqual([asset["description"] for asset in assets], ["图1", "图2", "图3", "图4"])
|