TelegramBotAdmin/tests/test_message_cleanup.py

from __future__ import annotations

import asyncio
from types import SimpleNamespace
from unittest.mock import AsyncMock

import pytest


class FakeMessage:
    def __init__(self, chat_id: int, message_id: int, *, text: str = "") -> None:
        self.chat = SimpleNamespace(
            id=chat_id,
            type=SimpleNamespace(value="supergroup" if chat_id < 0 else "private"),
        )
        self.id = message_id
        self.text = text
        self.caption = None
        self.delete_calls: list[bool] = []

    async def delete(self, *, revoke: bool = True) -> None:
        self.delete_calls.append(revoke)


class FakeClient:
    def __init__(self) -> None:
        self.next_message_id = 100

    async def send_message(
        self,
        chat_id: int,
        _text: str,
        *,
        reply_to_message_id: int | None = None,
    ) -> FakeMessage:
        self.next_message_id += 1
        return FakeMessage(chat_id, self.next_message_id)


@pytest.fixture
def cleanup(app_modules):
    module = app_modules.load("wbb.services.message_cleanup")
    module._scheduled_messages.clear()
    module._cleanup_tasks.clear()
    module._active_interactions.clear()
    yield module
    module._scheduled_messages.clear()
    module._cleanup_tasks.clear()
    module._active_interactions.clear()


async def test_group_cleanup_revokes_once_for_everyone(cleanup, monkeypatch):
    real_sleep = asyncio.sleep

    async def immediate_sleep(_delay: int) -> None:
        return None

    monkeypatch.setattr(cleanup.asyncio, "sleep", immediate_sleep)
    message = FakeMessage(-1001, 10, text="/points")

    cleanup.schedule_group_cleanup(message)
    cleanup.schedule_group_cleanup(message)
    await real_sleep(0)
    await real_sleep(0)

    assert message.delete_calls == [True]
    assert (-1001, 10) not in cleanup._scheduled_messages


async def test_private_messages_are_not_scheduled(cleanup):
    message = FakeMessage(1001, 11, text="/points")

    cleanup.schedule_group_cleanup(message)

    assert not cleanup.is_group_cleanup_scheduled(message)
    assert message.delete_calls == []


async def test_cleanup_task_is_referenced_until_deletion_finishes(
    cleanup, monkeypatch
):
    real_sleep = asyncio.sleep
    release = asyncio.Event()

    async def blocked_sleep(_delay: int) -> None:
        await release.wait()

    monkeypatch.setattr(cleanup.asyncio, "sleep", blocked_sleep)
    message = FakeMessage(-1001, 17, text="签到")

    cleanup.schedule_group_cleanup(message)
    await real_sleep(0)

    assert len(cleanup._cleanup_tasks) == 1
    release.set()
    await real_sleep(0)
    await real_sleep(0)

    assert message.delete_calls == [True]
    assert cleanup._cleanup_tasks == set()


async def test_command_context_tracks_only_responses_in_the_same_group(
    cleanup, monkeypatch
):
    scheduled: list[FakeMessage] = []
    monkeypatch.setattr(
        cleanup,
        "schedule_group_cleanup",
        lambda *messages, **_kwargs: scheduled.extend(messages),
    )
    client = FakeClient()
    cleanup.install_group_interaction_tracking(client)
    command = FakeMessage(-1001, 12, text="/points")

    assert cleanup.begin_group_message_interaction(command)
    group_response = await client.send_message(
        -1001, "积分余额", reply_to_message_id=command.id
    )
    await client.send_message(-1001, "抽奖主消息")
    await client.send_message(-1002, "其他群通知", reply_to_message_id=command.id)
    await client.send_message(1001, "私聊通知", reply_to_message_id=command.id)

    assert scheduled == [command, group_response]


async def test_plain_message_clears_previous_command_context(cleanup, monkeypatch):
    scheduled: list[FakeMessage] = []
    monkeypatch.setattr(
        cleanup,
        "schedule_group_cleanup",
        lambda *messages, **_kwargs: scheduled.extend(messages),
    )
    client = FakeClient()
    cleanup.install_group_interaction_tracking(client)
    command = FakeMessage(-1001, 13, text="/points")
    plain_message = FakeMessage(-1001, 14, text="大家好")

    assert cleanup.begin_group_message_interaction(command)
    assert not cleanup.begin_group_message_interaction(plain_message)
    await client.send_message(
        -1001, "普通自动回复", reply_to_message_id=plain_message.id
    )

    assert scheduled == [command]


async def test_checkin_button_is_treated_as_group_interaction(cleanup, monkeypatch):
    scheduled: list[FakeMessage] = []
    monkeypatch.setattr(
        cleanup,
        "schedule_group_cleanup",
        lambda *messages, **_kwargs: scheduled.extend(messages),
    )
    message = FakeMessage(-1001, 15, text="签到")

    assert cleanup.begin_group_message_interaction(message, menu_texts={"签到"})
    assert scheduled == [message]


async def test_callback_only_continues_an_already_ephemeral_menu(
    cleanup, monkeypatch
):
    scheduled: list[FakeMessage] = []
    monkeypatch.setattr(
        cleanup,
        "schedule_group_cleanup",
        lambda *messages, **_kwargs: scheduled.extend(messages),
    )
    client = FakeClient()
    cleanup.install_group_interaction_tracking(client)
    persistent_menu = FakeMessage(-1001, 16)
    query = SimpleNamespace(message=persistent_menu)

    assert not cleanup.begin_group_callback_interaction(query)
    await client.send_message(
        -1001, "不应自动删除", reply_to_message_id=persistent_menu.id
    )
    cleanup._scheduled_messages.add((-1001, 16))
    assert cleanup.begin_group_callback_interaction(query)
    callback_response = await client.send_message(
        -1001, "临时菜单回复", reply_to_message_id=persistent_menu.id
    )

    assert scheduled == [callback_response]


async def test_group_reply_explicitly_schedules_request_and_response(
    cleanup, monkeypatch
):
    scheduled: list[FakeMessage] = []
    monkeypatch.setattr(
        cleanup,
        "schedule_group_cleanup",
        lambda *messages, **_kwargs: scheduled.extend(messages),
    )
    request = FakeMessage(-1001, 18, text="/start")
    response = FakeMessage(-1001, 19, text="请私聊我查看详细功能。")
    request.reply = AsyncMock(return_value=response)

    result = await cleanup.reply_group_interaction(
        request,
        "请私聊我查看详细功能。",
        reply_markup=object(),
    )

    assert result is response
    request.reply.assert_awaited_once()
    assert scheduled == [request, response]


async def test_chinese_history_command_tracks_reply_with_configured_ttl(
    cleanup, monkeypatch
):
    scheduled: list[tuple[FakeMessage, int]] = []
    monkeypatch.setattr(
        cleanup,
        "schedule_group_cleanup",
        lambda *messages, delay_seconds=10: scheduled.extend(
            (message, delay_seconds) for message in messages
        ),
    )
    client = FakeClient()
    cleanup.install_group_interaction_tracking(client)
    command = FakeMessage(-1001, 20, text="流水")
    assert cleanup.begin_group_message_interaction(command, delay_seconds=12)
    response = await client.send_message(
        -1001, "最近积分流水", reply_to_message_id=command.id
    )
    assert scheduled == [(command, 12), (response, 12)]