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="主插图", source_span={"paragraph_index": 0}, 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="主插图", source_span={"paragraph_index": 0}, 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()