66 lines
1.7 KiB
Python
66 lines
1.7 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from app.features.memoir.ingest_scheduler import MemoirIngestScheduler
|
||
|
|
|
||
|
|
|
||
|
|
class _FakeRunner:
|
||
|
|
def __init__(self) -> None:
|
||
|
|
self.queued: list[tuple[str, str, int]] = []
|
||
|
|
self.flushed: list[tuple[str, list[str]]] = []
|
||
|
|
|
||
|
|
async def queue_message(
|
||
|
|
self,
|
||
|
|
user_id: str,
|
||
|
|
segment_id: str,
|
||
|
|
*,
|
||
|
|
text_char_count: int = 0,
|
||
|
|
) -> None:
|
||
|
|
self.queued.append((user_id, segment_id, text_char_count))
|
||
|
|
|
||
|
|
async def flush_pending(
|
||
|
|
self,
|
||
|
|
user_id: str,
|
||
|
|
*,
|
||
|
|
extra_segment_ids: list[str] | None = None,
|
||
|
|
) -> str:
|
||
|
|
self.flushed.append((user_id, list(extra_segment_ids or [])))
|
||
|
|
return "task-1"
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_queue_segment_returns_visible_phase_plan() -> None:
|
||
|
|
runner = _FakeRunner()
|
||
|
|
scheduler = MemoirIngestScheduler(runner=runner)
|
||
|
|
|
||
|
|
plan = await scheduler.queue_segment(
|
||
|
|
"user-1",
|
||
|
|
"seg-1",
|
||
|
|
text_char_count=42,
|
||
|
|
trigger="evaluation_replay",
|
||
|
|
)
|
||
|
|
|
||
|
|
assert runner.queued == [("user-1", "seg-1", 42)]
|
||
|
|
assert plan.user_id == "user-1"
|
||
|
|
assert plan.segment_ids == ("seg-1",)
|
||
|
|
assert plan.trigger == "evaluation_replay"
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_flush_pending_returns_plan_and_task_id() -> None:
|
||
|
|
runner = _FakeRunner()
|
||
|
|
scheduler = MemoirIngestScheduler(runner=runner)
|
||
|
|
|
||
|
|
plan, task_id = await scheduler.flush_pending(
|
||
|
|
"user-1",
|
||
|
|
extra_segment_ids=["seg-1", "seg-2"],
|
||
|
|
trigger="conversation_end",
|
||
|
|
)
|
||
|
|
|
||
|
|
assert runner.flushed == [("user-1", ["seg-1", "seg-2"])]
|
||
|
|
assert task_id == "task-1"
|
||
|
|
assert plan.user_id == "user-1"
|
||
|
|
assert plan.segment_ids == ("seg-1", "seg-2")
|
||
|
|
assert plan.trigger == "conversation_end"
|