152 lines
5.1 KiB
Python
152 lines
5.1 KiB
Python
import unittest
|
|
from contextlib import contextmanager
|
|
from io import BytesIO
|
|
from types import SimpleNamespace
|
|
from unittest.mock import Mock, patch
|
|
|
|
from PIL import Image
|
|
|
|
from app.ports.image_gen import ImageResult, TaskStatus
|
|
from app.tasks.story_image_tasks import generate_story_image
|
|
|
|
|
|
def _mock_db_cm(db):
|
|
@contextmanager
|
|
def _cm():
|
|
yield db
|
|
|
|
return _cm()
|
|
|
|
|
|
def _png_bytes() -> bytes:
|
|
buf = BytesIO()
|
|
Image.new("RGB", (1, 1), color="white").save(buf, format="PNG")
|
|
return buf.getvalue()
|
|
|
|
|
|
class _FakeUUID:
|
|
def __init__(self, value: str):
|
|
self.hex = value
|
|
self._value = value
|
|
|
|
def __str__(self) -> str:
|
|
return self._value
|
|
|
|
|
|
class GenerateStoryImageTaskTest(unittest.TestCase):
|
|
@patch("app.tasks.story_image_tasks.release_redis_lock")
|
|
@patch(
|
|
"app.tasks.story_image_tasks.acquire_redis_lock",
|
|
return_value=SimpleNamespace(key="lock:story-image:story-1"),
|
|
)
|
|
@patch("app.tasks.story_image_tasks._claim_story_image_intent_sync")
|
|
@patch("app.tasks.story_image_tasks.get_sync_db")
|
|
@patch("app.tasks.story_image_tasks.TencentCosStorageService")
|
|
@patch("app.tasks.story_image_tasks.get_image_generator")
|
|
@patch("app.features.memoir.memoir_images.settings.MemoirImageSettings.from_env")
|
|
@patch("app.tasks.story_image_tasks.uuid.uuid4")
|
|
def test_generate_story_image_resumes_processing_intent_and_backfills_markdown(
|
|
self,
|
|
uuid4_mock,
|
|
settings_from_env,
|
|
get_image_generator_mock,
|
|
storage_cls,
|
|
get_sync_db_mock,
|
|
claim_intent_mock,
|
|
acquire_lock_mock,
|
|
release_lock_mock,
|
|
):
|
|
uuid4_mock.side_effect = [
|
|
_FakeUUID("claim-token"),
|
|
_FakeUUID("asset-uuid"),
|
|
_FakeUUID("version-uuid"),
|
|
]
|
|
settings_from_env.return_value = SimpleNamespace(
|
|
provider="liblib",
|
|
default_style="watercolor",
|
|
default_size="1024x1024",
|
|
)
|
|
|
|
intent = SimpleNamespace(
|
|
id="intent-1",
|
|
prompt_brief="院子里的藤椅",
|
|
style_profile="watercolor",
|
|
story_version_id="ver-1",
|
|
caption="主插图",
|
|
status="processing",
|
|
)
|
|
story = SimpleNamespace(
|
|
id="story-1",
|
|
user_id="user-1",
|
|
title="童年的院子",
|
|
stage="childhood",
|
|
)
|
|
db_claim = Mock()
|
|
claim_intent_mock.return_value = (intent, story)
|
|
|
|
intent_db = SimpleNamespace(
|
|
id="intent-1",
|
|
story_version_id="ver-1",
|
|
caption="主插图",
|
|
prompt_brief="院子里的藤椅",
|
|
status="processing",
|
|
style_profile="watercolor",
|
|
claim_token="claim-token",
|
|
asset_id=None,
|
|
error=None,
|
|
updated_at=None,
|
|
)
|
|
story_db = SimpleNamespace(
|
|
id="story-1",
|
|
current_version_id="ver-1",
|
|
canonical_markdown="第一段\n\n第二段",
|
|
)
|
|
version_db = SimpleNamespace(id="ver-1", markdown_snapshot="第一段\n\n第二段")
|
|
version_max_result = Mock()
|
|
version_max_result.scalar.return_value = 1
|
|
db_persist = Mock()
|
|
db_persist.get.side_effect = [intent_db, story_db, version_db]
|
|
db_persist.execute.return_value = version_max_result
|
|
|
|
get_sync_db_mock.side_effect = [_mock_db_cm(db_claim), _mock_db_cm(db_persist)]
|
|
|
|
generator = get_image_generator_mock.return_value
|
|
generator.generate.return_value = ImageResult(
|
|
status=TaskStatus.COMPLETED,
|
|
task_id="task-1",
|
|
image_url="https://provider.example.com/story.png",
|
|
)
|
|
generator.download_image.return_value = _png_bytes()
|
|
storage_cls.from_env.return_value.upload_bytes.return_value = (
|
|
"https://cos.example.com/stories/u1/s1.png"
|
|
)
|
|
|
|
result = generate_story_image.run("story-1")
|
|
|
|
self.assertEqual(result["status"], "success")
|
|
self.assertEqual(intent_db.status, "completed")
|
|
self.assertIsNotNone(intent_db.asset_id)
|
|
self.assertNotEqual(story_db.current_version_id, "ver-1")
|
|
self.assertIn("asset://", story_db.canonical_markdown)
|
|
generator.generate.assert_called_once()
|
|
storage_cls.from_env.return_value.upload_bytes.assert_called_once()
|
|
claim_intent_mock.assert_called_once()
|
|
acquire_lock_mock.assert_called_once()
|
|
release_lock_mock.assert_called_once()
|
|
|
|
@patch("app.tasks.story_image_tasks.acquire_redis_lock", return_value=None)
|
|
@patch("app.tasks.story_image_tasks.get_sync_db")
|
|
@patch("app.tasks.story_image_tasks.get_image_generator")
|
|
def test_generate_story_image_skips_when_lock_is_already_held(
|
|
self,
|
|
get_image_generator_mock,
|
|
get_sync_db_mock,
|
|
acquire_lock_mock,
|
|
):
|
|
result = generate_story_image.run("story-1")
|
|
|
|
self.assertEqual(result, {"status": "locked"})
|
|
get_sync_db_mock.assert_not_called()
|
|
get_image_generator_mock.assert_not_called()
|
|
acquire_lock_mock.assert_called_once()
|