life-echo/api/tests/test_process_memoir_segments_image_enqueue.py

import unittest
from contextlib import contextmanager
from types import SimpleNamespace
from unittest.mock import Mock, patch

from app.tasks.memoir_tasks import MemoirImageSettings, process_memoir_segments


def _mock_get_sync_db(db):
    @contextmanager
    def _cm():
        yield db

    return _cm()  # 返回 context manager 实例，供 with 使用


class ProcessMemoirSegmentsImageEnqueueTest(unittest.TestCase):
    @patch(
        "app.features.memory.repo.retrieve_evidence_sync",
        return_value={
            "relevant_chunks": [],
            "relevant_summaries": [],
            "relevant_facts": [],
            "timeline_hints": [],
            "relevant_stories": [],
        },
    )
    @patch("app.features.memory.service.ingest_transcript_sync")
    @patch("app.tasks.memoir_tasks._update_task_status_sync")
    @patch("app.tasks.memoir_tasks._release_chapter_lock")
    @patch("app.tasks.memoir_tasks._acquire_chapter_lock", return_value=True)
    @patch("app.tasks.memoir_tasks._update_slot_sync")
    @patch(
        "app.agents.memoir.orchestrator.ClassificationAgent.classify",
        return_value="childhood",
    )
    @patch("app.tasks.memoir_tasks._get_or_create_state_sync")
    @patch("app.tasks.memoir_tasks._get_llm")
    @patch("app.tasks.chapter_cover_tasks.generate_chapter_cover.delay")
    @patch("app.tasks.memoir_tasks.get_sync_db")
    @patch("app.tasks.memoir_tasks.MemoirImageSettings.from_env")
    def test_process_memoir_segments_parses_markdown_wrapped_state_extraction_json(
        self,
        settings_from_env,
        get_sync_db_mock,
        delay_mock,
        get_llm_mock,
        get_state_mock,
        _classify_mock,
        update_slot_mock,
        _acquire_lock_mock,
        _release_lock_mock,
        _update_status_mock,
        ingest_mock,
        retrieve_mock,
    ):
        settings_from_env.return_value = MemoirImageSettings(
            enabled=True,
            max_per_chapter=2,
            provider="liblib",
            default_style="watercolor",
            default_size="1024x1024",
            poll_interval_seconds=3,
            max_attempts=20,
            liblib_template_uuid="tpl-uuid",
        )
        get_state_mock.return_value = SimpleNamespace(
            current_stage="childhood", slots={}
        )
        update_slot_mock.return_value = SimpleNamespace(
            current_stage="childhood", slots={}
        )
        llm = Mock()
        bound_llm = Mock()
        bound_llm.invoke.side_effect = [
            SimpleNamespace(
                content="""```json
{
  "detected_stage": "childhood",
  "slots": {
    "family_memory": "外婆总在门口等我"
  }
}
```"""
            ),
            SimpleNamespace(
                content='{"paragraphs":[{"content":"新的章节正文","image_description":"南方小镇的青石板路"}]}'
            ),
        ]
        llm.bind.return_value = bound_llm
        llm.invoke.side_effect = [
            SimpleNamespace(content="childhood"),
            SimpleNamespace(content="童年的门前"),
        ]
        get_llm_mock.return_value = llm

        segment = SimpleNamespace(
            id="segment-1",
            transcript_text="那条路我一直记得。\n\n{{{{IMAGE:南方小镇的青石板路}}}}",
            processed=False,
        )

        segments_result = Mock()
        segments_result.scalars.return_value.all.return_value = [segment]

        unique_result = Mock()
        unique_result.scalar_one_or_none.return_value = None
        chapter_result = Mock()
        chapter_result.unique.return_value = unique_result

        book_result = Mock()
        book_result.scalar_one_or_none.return_value = None

        empty_sections_result = Mock()
        empty_sections_result.scalars.return_value.all.return_value = []

        version_count_result = Mock()
        version_count_result.scalar.return_value = 0

        db = Mock()
        db.execute.side_effect = [
            segments_result,
            chapter_result,
            empty_sections_result,  # _save_narrative_to_sections 内查询 ChapterSection
            version_count_result,  # ensure_chapter_markdown_and_version_sync 内 count
            book_result,
        ]
        db.get.return_value = None
        get_sync_db_mock.return_value = _mock_get_sync_db(db)

        events: list[str] = []
        db.commit.side_effect = lambda: events.append("commit")
        delay_mock.side_effect = lambda chapter_id: events.append(f"delay:{chapter_id}")

        task_self = SimpleNamespace(
            request=SimpleNamespace(id="task-1"),
            retry=Mock(side_effect=AssertionError("retry should not be called")),
        )
        process_memoir_segments.run.__func__(task_self, "user-1", ["segment-1"])

        update_slot_mock.assert_called_once_with(
            "user-1",
            "childhood",
            "family_memory",
            "外婆总在门口等我",
            ["segment-1"],
            db,
        )
        self.assertIn("commit", events)
        delay_events = [event for event in events if event.startswith("delay:")]
        self.assertEqual(len(delay_events), 1)
        self.assertGreater(events.index(delay_events[0]), events.index("commit"))

    @patch(
        "app.features.memory.repo.retrieve_evidence_sync",
        return_value={
            "relevant_chunks": [],
            "relevant_summaries": [],
            "relevant_facts": [],
            "timeline_hints": [],
            "relevant_stories": [],
        },
    )
    @patch("app.features.memory.service.ingest_transcript_sync")
    @patch("app.tasks.memoir_tasks._update_task_status_sync")
    @patch("app.tasks.memoir_tasks._release_chapter_lock")
    @patch("app.tasks.memoir_tasks._acquire_chapter_lock", return_value=True)
    @patch(
        "app.agents.memoir.orchestrator.ClassificationAgent.classify",
        return_value="childhood",
    )
    @patch("app.tasks.memoir_tasks._get_or_create_state_sync")
    @patch("app.tasks.memoir_tasks._get_llm", return_value=None)
    @patch("app.tasks.chapter_cover_tasks.generate_chapter_cover.delay")
    @patch("app.tasks.memoir_tasks.get_sync_db")
    @patch("app.tasks.memoir_tasks.MemoirImageSettings.from_env")
    def test_process_memoir_segments_does_not_enqueue_image_jobs_when_feature_disabled(
        self,
        settings_from_env,
        get_sync_db_mock,
        delay_mock,
        _get_llm,
        get_state_mock,
        _classify_mock,
        _acquire_lock_mock,
        _release_lock_mock,
        _update_status_mock,
        ingest_mock,
        retrieve_mock,
    ):
        settings_from_env.return_value = MemoirImageSettings(
            enabled=False,
            max_per_chapter=2,
            provider="liblib",
            default_style="watercolor",
            default_size="1024x1024",
            poll_interval_seconds=3,
            max_attempts=20,
            liblib_template_uuid="tpl-uuid",
        )
        get_state_mock.return_value = SimpleNamespace(
            current_stage="childhood", slots={}
        )

        segment = SimpleNamespace(
            id="segment-1",
            transcript_text="那条路我一直记得。\n\n{{{{IMAGE:南方小镇的青石板路}}}}",
            processed=False,
        )

        segments_result = Mock()
        segments_result.scalars.return_value.all.return_value = [segment]

        unique_result = Mock()
        unique_result.scalar_one_or_none.return_value = None
        chapter_result = Mock()
        chapter_result.unique.return_value = unique_result

        book_result = Mock()
        book_result.scalar_one_or_none.return_value = None

        empty_sections_result = Mock()
        empty_sections_result.scalars.return_value.all.return_value = []

        version_count_result = Mock()
        version_count_result.scalar.return_value = 0

        db = Mock()
        db.execute.side_effect = [
            segments_result,
            chapter_result,
            empty_sections_result,  # _save_narrative_to_sections 内查询 ChapterSection
            version_count_result,  # ensure_chapter_markdown_and_version_sync 内 count
            book_result,
        ]
        db.get.return_value = None
        get_sync_db_mock.return_value = _mock_get_sync_db(db)

        task_self = SimpleNamespace(
            request=SimpleNamespace(id="task-1"),
            retry=Mock(side_effect=AssertionError("retry should not be called")),
        )
        process_memoir_segments.run.__func__(task_self, "user-1", ["segment-1"])

        delay_mock.assert_not_called()