steward_mirror/tests/test_bot.py

"""Tests for steward.bot.telegram (handler logic)."""

from unittest.mock import AsyncMock, MagicMock

import pytest
from telegram import Chat, Message, Update, User
from telegram.ext import CallbackContext

from steward.bot.core import ConversationService
from steward.bot.telegram import (
    _is_allowed,
    _is_chat_enabled,
    clear_handler,
    message_handler,
    start_handler,
)
from steward.bot.thread_key import ThreadKey
from steward.config import Settings
from steward.llm.client import LLMClient
from steward.memory.thread_store import ThreadMemoryStore
from tests.conftest import make_settings


def _make_settings(**kwargs) -> Settings:
    """Wrapper for test compatibility."""
    defaults = dict(
        telegram_bot_token="test-token",
        openai_api_key="test-key",
    )
    defaults.update(kwargs)
    return make_settings(**defaults)


def _make_update(
    user_id: int = 12345,
    text: str = "hello",
    chat_id: int | None = None,
    thread_id: int | None = None,
    chat_type: str = "private",
) -> Update:
    user = MagicMock(spec=User)
    user.id = user_id

    message = MagicMock(spec=Message)
    message.text = text
    message.reply_text = AsyncMock()
    message.reply_poll = AsyncMock()
    message.message_thread_id = thread_id

    chat = MagicMock(spec=Chat)
    chat.id = chat_id if chat_id is not None else user_id
    chat.type = chat_type

    update = MagicMock(spec=Update)
    update.effective_user = user
    update.effective_chat = chat
    update.message = message
    return update


def _make_context(
    settings: Settings,
    llm: LLMClient | None = None,
    store: ThreadMemoryStore | None = None,
) -> CallbackContext:  # type: ignore[type-arg]
    ctx = MagicMock(spec=CallbackContext)
    if store is None:
        mock_store = MagicMock(spec=ThreadMemoryStore)
        mock_store.search.return_value = []
        store = mock_store
    if llm is None:
        llm = MagicMock(spec=LLMClient)
    service = ConversationService(settings, llm, store)
    ctx.bot_data = {
        "settings": settings,
        "service": service,
        "llm": llm,
    }
    ctx.args = []
    return ctx


class TestIsAllowed:
    def test_no_allowlist_allows_everyone(self):
        settings = _make_settings(telegram_allowed_user_ids=[])
        assert _is_allowed(999, settings) is True

    def test_allowlist_accepts_known_user(self):
        settings = _make_settings(telegram_allowed_user_ids=[1, 2, 3])
        assert _is_allowed(2, settings) is True

    def test_allowlist_rejects_unknown_user(self):
        settings = _make_settings(telegram_allowed_user_ids=[1, 2, 3])
        assert _is_allowed(999, settings) is False


class TestIsChatEnabled:
    def test_private_chat_ignores_group_allowlist(self):
        settings = _make_settings(telegram_group_ids=[-100123])
        update = _make_update(user_id=123, chat_type="private")

        assert _is_chat_enabled(update.effective_chat, settings) is True

    def test_group_chat_accepts_configured_group(self):
        settings = _make_settings(telegram_group_ids=[-5308306472])
        update = _make_update(chat_id=-5308306472, chat_type="group")

        assert _is_chat_enabled(update.effective_chat, settings) is True

    def test_group_chat_rejects_unconfigured_group(self):
        settings = _make_settings(telegram_group_ids=[-1005308306472])
        update = _make_update(chat_id=-5308306472, chat_type="group")

        assert _is_chat_enabled(update.effective_chat, settings) is False


@pytest.mark.asyncio
async def test_start_handler_replies(monkeypatch):
    settings = _make_settings()
    update = _make_update()
    ctx = _make_context(settings)

    await start_handler(update, ctx)

    update.message.reply_text.assert_awaited_once()
    text = update.message.reply_text.call_args.args[0]
    assert "Steward" in text


@pytest.mark.asyncio
async def test_start_handler_ignores_disallowed_user():
    settings = _make_settings(telegram_allowed_user_ids=[9999])
    update = _make_update(user_id=1111)
    ctx = _make_context(settings)

    await start_handler(update, ctx)

    update.message.reply_text.assert_not_awaited()


@pytest.mark.asyncio
async def test_clear_handler_clears_user_history():
    """clear_handler without a thread clears per-user history."""
    settings = _make_settings()
    user_id = 42
    ctx = _make_context(settings)
    service: ConversationService = ctx.bot_data["service"]
    key = ThreadKey(platform="telegram", scope=str(user_id))
    service._histories[key] = [{"role": "user", "content": "old msg"}]

    update = _make_update(user_id=user_id)  # no thread_id
    await clear_handler(update, ctx)

    assert not service.has_history(key)
    update.message.reply_text.assert_awaited_once()


@pytest.mark.asyncio
async def test_clear_handler_clears_thread_history():
    """clear_handler inside a thread clears that thread's history."""
    settings = _make_settings()
    ctx = _make_context(settings)
    service: ConversationService = ctx.bot_data["service"]
    key = ThreadKey(platform="telegram", scope="100", thread="7")
    service._histories[key] = [{"role": "user", "content": "thread msg"}]

    update = _make_update(user_id=1, chat_id=100, thread_id=7)
    await clear_handler(update, ctx)

    assert not service.has_history(key)
    update.message.reply_text.assert_awaited_once()


@pytest.mark.asyncio
async def test_message_handler_calls_llm_and_replies():
    settings = _make_settings()
    mock_llm = MagicMock(spec=LLMClient)
    mock_llm.chat = AsyncMock(return_value="LLM response")

    update = _make_update(user_id=77, text="What is the weather?")
    ctx = _make_context(settings, llm=mock_llm)
    service: ConversationService = ctx.bot_data["service"]
    key = ThreadKey(platform="telegram", scope="77")

    await message_handler(update, ctx)

    mock_llm.chat.assert_awaited_once()
    update.message.reply_text.assert_awaited_once()
    history = service._histories[key]
    assert len(history) == 2
    assert history[0]["role"] == "user"
    assert history[1]["role"] == "assistant"


@pytest.mark.asyncio
async def test_message_handler_sends_multiple_reply_messages():
    settings = _make_settings()
    mock_llm = MagicMock(spec=LLMClient)
    mock_llm.chat = AsyncMock(
        return_value="[MESSAGE]\nFirst thought.\n[MESSAGE]\nWhat outcome do you want?"
    )

    update = _make_update(user_id=78, text="Can you help with this vague thing?")
    ctx = _make_context(settings, llm=mock_llm)
    service: ConversationService = ctx.bot_data["service"]
    key = ThreadKey(platform="telegram", scope="78")

    await message_handler(update, ctx)

    assert update.message.reply_text.await_count == 2
    assert update.message.reply_text.await_args_list[0].args[0] == "First thought."
    assert update.message.reply_text.await_args_list[1].args[0] == "What outcome do you want?"
    update.message.reply_poll.assert_not_awaited()
    assert service._histories[key][1]["content"] == "First thought.\n\nWhat outcome do you want?"


@pytest.mark.asyncio
async def test_message_handler_sends_native_poll_from_directive():
    settings = _make_settings()
    mock_llm = MagicMock(spec=LLMClient)
    mock_llm.chat = AsyncMock(
        return_value=(
            "[MESSAGE]\nI can turn that into a vote.\n"
            "[POLL]\n"
            "question: Which implementation should we choose?\n"
            "- Minimal change\n"
            "- Full redesign\n"
            "[/POLL]\n"
            "[MESSAGE]\nI'll use the winning option."
        )
    )

    update = _make_update(user_id=79, text="Should we do a minimal change or redesign?")
    ctx = _make_context(settings, llm=mock_llm)
    service: ConversationService = ctx.bot_data["service"]
    key = ThreadKey(platform="telegram", scope="79")

    await message_handler(update, ctx)

    assert update.message.reply_text.await_count == 2
    update.message.reply_poll.assert_awaited_once_with(
        question="Which implementation should we choose?",
        options=["Minimal change", "Full redesign"],
        is_anonymous=False,
    )
    assert service._histories[key][1]["content"] == (
        "I can turn that into a vote.\n\n"
        "Poll: Which implementation should we choose? (Minimal change; Full redesign)\n\n"
        "I'll use the winning option."
    )


@pytest.mark.asyncio
async def test_message_handler_non_thread_trims_history():
    """Non-thread history is trimmed to _MAX_HISTORY turns."""
    from steward.bot.core import _MAX_HISTORY

    settings = _make_settings()
    mock_llm = MagicMock(spec=LLMClient)
    mock_llm.chat = AsyncMock(return_value="reply")

    user_id = 200
    ctx = _make_context(settings, llm=mock_llm)
    service: ConversationService = ctx.bot_data["service"]
    key = ThreadKey(platform="telegram", scope=str(user_id))
    # Pre-fill exactly at the limit
    service._histories[key] = [
        {"role": "user" if i % 2 == 0 else "assistant", "content": f"msg{i}"}
        for i in range(_MAX_HISTORY * 2)
    ]

    update = _make_update(user_id=user_id, text="new question")
    await message_handler(update, ctx)

    # After appending 2 new messages, should trim back to _MAX_HISTORY * 2
    assert len(service._histories[key]) == _MAX_HISTORY * 2