from __future__ import annotations from datetime import datetime, timezone from types import SimpleNamespace from uuid import uuid4 import pytest from models.agent_chat_message import AgentChatMessageRole from v1.agent.repository import AgentRepository class _ExecuteResult: def __init__(self, value: object) -> None: self._value = value def scalar_one_or_none(self) -> object: return self._value class _FakeSession: def __init__(self, session_row: object) -> None: self.session_row = session_row self.added: list[object] = [] self.flushed = False async def execute(self, stmt): # noqa: ANN001 del stmt return _ExecuteResult(self.session_row) def add(self, obj: object) -> None: self.added.append(obj) async def flush(self) -> None: self.flushed = True @pytest.mark.asyncio async def test_snapshot_message_returns_raw_db_columns() -> None: repository = AgentRepository(session=SimpleNamespace()) # type: ignore[arg-type] now = datetime.now(timezone.utc) message = SimpleNamespace( id=uuid4(), session_id=uuid4(), seq=7, role=AgentChatMessageRole.TOOL, content='{"offloaded":true}', metadata_json={"tool_call_id": "call-1"}, created_at=now, ) payload = await repository._to_snapshot_message(message) # type: ignore[arg-type] assert payload["seq"] == 7 assert payload["role"] == "tool" assert payload["content"] == '{"offloaded":true}' assert payload["metadata"] == {"tool_call_id": "call-1"} assert "timestamp" in payload @pytest.mark.asyncio async def test_persist_user_message_sets_session_title_when_empty() -> None: session_id = str(uuid4()) session_row = SimpleNamespace( message_count=0, title=None, last_activity_at=datetime.now(timezone.utc), ) fake_session = _FakeSession(session_row) repository = AgentRepository(session=fake_session) # type: ignore[arg-type] await repository.persist_user_message( session_id=session_id, run_id="run-1", content_text=" 请帮我安排明天下午开会 ", metadata=None, ) assert session_row.title == "请帮我安排明天下午开会" assert session_row.message_count == 1 assert fake_session.flushed is True @pytest.mark.asyncio async def test_persist_user_message_keeps_existing_session_title() -> None: session_id = str(uuid4()) session_row = SimpleNamespace( message_count=1, title="已有标题", last_activity_at=datetime.now(timezone.utc), ) fake_session = _FakeSession(session_row) repository = AgentRepository(session=fake_session) # type: ignore[arg-type] await repository.persist_user_message( session_id=session_id, run_id="run-2", content_text="新的消息内容", metadata=None, ) assert session_row.title == "已有标题" assert session_row.message_count == 2