data-query-agent/tests/test_direct_reply.py

from types import SimpleNamespace

import pytest

from data_query_agent.models import IncomingMessage, QueryDecision, SkillParameters
from data_query_agent.service import DataQueryService


def _parameters() -> SkillParameters:
    return SkillParameters(**{name: None for name in SkillParameters.model_fields})


class FakeState:
    def __init__(self) -> None:
        self.saved_thread: tuple[str, str] | None = None

    def get_conversation(self, key: str) -> SimpleNamespace:
        return SimpleNamespace(thread_id="old-thread", session_id="session-1")

    def set_thread(self, key: str, thread_id: str) -> None:
        self.saved_thread = (key, thread_id)

    def create_run(self, *args: object, **kwargs: object) -> None:
        raise AssertionError("direct_reply must not create a query run")


class FakeCodex:
    async def plan(self, thread_id: str | None, question: str) -> tuple[str, QueryDecision]:
        assert thread_id == "old-thread"
        assert question == "分析下刚才结果"
        return "new-thread", QueryDecision(
            status="ready",
            reply="实验组整体效率小幅走弱，主要来自推荐场景回流下降。",
            title="数据助手回复",
            selected_skill="odps-product-efficiency-report",
            execution_mode="direct_reply",
            sql=None,
            parameters=_parameters(),
            assumptions=["仅基于本线程已有查询结果"],
        )


class FakeFeishu:
    def __init__(self) -> None:
        self.replies: list[tuple[str, str]] = []

    async def reply_text(self, message_id: str, text: str) -> None:
        self.replies.append((message_id, text))


@pytest.mark.asyncio
async def test_direct_reply_uses_thread_context_without_starting_query() -> None:
    service = object.__new__(DataQueryService)
    service.state = FakeState()
    service.codex = FakeCodex()
    service.feishu = FakeFeishu()
    message = IncomingMessage(
        message_id="m1",
        chat_id="c1",
        chat_type="group",
        sender_open_id="u1",
        text="分析下刚才结果",
        mentioned_bot=True,
    )

    await service._handle_query(message)

    assert service.state.saved_thread == (message.conversation_key, "new-thread")
    assert service.feishu.replies == [
        ("m1", "实验组整体效率小幅走弱，主要来自推荐场景回流下降。")
    ]