Fix memoir image prompt parsing and host allowlist

This commit is contained in:
Kevin
2026-03-11 13:18:20 +08:00
parent 1f98b8bfd6
commit 32954d4b3f
5 changed files with 100 additions and 3 deletions

View File

@@ -63,8 +63,12 @@ class GenerateChapterImagesTaskTest(unittest.TestCase):
@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,
@@ -118,8 +122,12 @@ class GenerateChapterImagesTaskTest(unittest.TestCase):
@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,
@@ -160,7 +168,9 @@ class GenerateChapterImagesTaskTest(unittest.TestCase):
"status": "completed",
"image_url": "https://provider.example.com/1.png",
}
provider_inst.download_image.return_value = b"png-bytes"
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"
@@ -235,8 +245,12 @@ class GenerateChapterImagesTaskTest(unittest.TestCase):
@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,
@@ -295,8 +309,12 @@ class GenerateChapterImagesTaskTest(unittest.TestCase):
@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,