social-app/backend/tests/unit/v1/agent/test_repository.py

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