life-echo/api/tests/test_story_image_tasks.py

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()