240 lines
8.1 KiB
Python
240 lines
8.1 KiB
Python
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 使用
|
|
|
|
|
|
def _fake_chapter_for_pipeline():
|
|
return SimpleNamespace(
|
|
id="chapter-1",
|
|
canonical_markdown="# 标题\n\n正文若干字。",
|
|
cover_asset_id=None,
|
|
images=[],
|
|
)
|
|
|
|
|
|
class ProcessMemoirSegmentsImageEnqueueTest(unittest.TestCase):
|
|
@patch("app.tasks.chapter_compose_tasks.recompose_chapters_for_story.delay")
|
|
@patch("app.tasks.story_image_tasks.generate_story_image.delay")
|
|
@patch("app.tasks.memoir_tasks.run_story_pipeline_for_category_batch")
|
|
@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_enqueue.try_enqueue_generate_chapter_cover")
|
|
@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,
|
|
try_enqueue_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,
|
|
mock_pipeline,
|
|
_delay_story_image,
|
|
_delay_recompose,
|
|
):
|
|
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={}
|
|
)
|
|
mock_pipeline.return_value = (_fake_chapter_for_pipeline(), True, {"story-1"})
|
|
llm = Mock()
|
|
bound_llm = Mock()
|
|
bound_llm.invoke.side_effect = [
|
|
SimpleNamespace(
|
|
content="""```json
|
|
{
|
|
"detected_stage": "childhood",
|
|
"slots": {
|
|
"family_memory": "外婆总在门口等我"
|
|
}
|
|
}
|
|
```"""
|
|
),
|
|
]
|
|
llm.bind.return_value = bound_llm
|
|
llm.invoke.side_effect = [
|
|
SimpleNamespace(content="childhood"),
|
|
]
|
|
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]
|
|
|
|
book_result = Mock()
|
|
book_result.scalar_one_or_none.return_value = None
|
|
|
|
db = Mock()
|
|
db.execute.side_effect = [
|
|
segments_result,
|
|
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")
|
|
try_enqueue_mock.side_effect = lambda chapter_id, source="pipeline": (
|
|
events.append(f"enqueue:{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,
|
|
)
|
|
mock_pipeline.assert_called_once()
|
|
self.assertIn("commit", events)
|
|
enqueue_events = [event for event in events if event.startswith("enqueue:")]
|
|
self.assertEqual(len(enqueue_events), 1)
|
|
self.assertGreater(events.index(enqueue_events[0]), events.index("commit"))
|
|
|
|
@patch("app.tasks.chapter_compose_tasks.recompose_chapters_for_story.delay")
|
|
@patch("app.tasks.story_image_tasks.generate_story_image.delay")
|
|
@patch("app.tasks.memoir_tasks.run_story_pipeline_for_category_batch")
|
|
@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_enqueue.try_enqueue_generate_chapter_cover")
|
|
@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,
|
|
try_enqueue_mock,
|
|
_get_llm,
|
|
get_state_mock,
|
|
_classify_mock,
|
|
_acquire_lock_mock,
|
|
_release_lock_mock,
|
|
_update_status_mock,
|
|
ingest_mock,
|
|
retrieve_mock,
|
|
mock_pipeline,
|
|
_delay_story_image,
|
|
_delay_recompose,
|
|
):
|
|
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={}
|
|
)
|
|
mock_pipeline.return_value = (_fake_chapter_for_pipeline(), True, set())
|
|
|
|
segment = SimpleNamespace(
|
|
id="segment-1",
|
|
transcript_text="那条路我一直记得。\n\n{{{{IMAGE:南方小镇的青石板路}}}}",
|
|
processed=False,
|
|
)
|
|
|
|
segments_result = Mock()
|
|
segments_result.scalars.return_value.all.return_value = [segment]
|
|
|
|
book_result = Mock()
|
|
book_result.scalar_one_or_none.return_value = None
|
|
|
|
db = Mock()
|
|
db.execute.side_effect = [
|
|
segments_result,
|
|
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"])
|
|
|
|
try_enqueue_mock.assert_not_called()
|