TelegramBotAdmin/tests/test_captcha_verification.py

from __future__ import annotations

import asyncio
import sys
from types import ModuleType, SimpleNamespace
from unittest.mock import AsyncMock

import pytest
from pyrogram.enums import ChatMemberStatus


@pytest.fixture
async def greetings(app_modules, monkeypatch):
    app_modules.wbb.WELCOME_DELAY_KICK_SEC = 300
    permissions = ModuleType("wbb.core.decorators.permissions")
    permissions.adminsOnly = lambda _permission: lambda function: function
    notes = ModuleType("wbb.modules.notes")
    notes.extract_urls = lambda _markup: []
    monkeypatch.setitem(sys.modules, permissions.__name__, permissions)
    monkeypatch.setitem(sys.modules, notes.__name__, notes)
    module = app_modules.load("wbb.modules.greetings")
    await asyncio.sleep(0)
    return module


async def test_captcha_defaults_off_and_admin_toggle_matches_bot(app_modules, monkeypatch):
    dbfunctions = app_modules.load("wbb.utils.dbfunctions")
    service = app_modules.load("wbb.services.chat_management")
    monkeypatch.setattr(service, "ensure_permission", AsyncMock())

    assert await dbfunctions.is_captcha_on(-100) is False
    assert (await service.get_automation_settings(-100))["captcha_enabled"] is False
    await app_modules.wbb.db.managed_chat_settings.insert_one(
        {"bot_id": "primary", "chat_id": -200, "welcome": {"enabled": True}}
    )
    assert await dbfunctions.is_captcha_on(-200) is False
    assert (await service.get_automation_settings(-200))["captcha_enabled"] is False

    await service.apply_automation_settings(-100, {"captcha_enabled": True})
    assert await dbfunctions.is_captcha_on(-100) is True
    await dbfunctions.captcha_off(-100)  # The /captcha command uses the same helper.
    assert await dbfunctions.is_captcha_on(-100) is False
    assert (await service.get_automation_settings(-100))["captcha_enabled"] is False
    await dbfunctions.captcha_on(-100)
    assert (await service.get_automation_settings(-100))["captcha_enabled"] is True


async def test_solved_captcha_is_saved_for_each_member(app_modules):
    dbfunctions = app_modules.load("wbb.utils.dbfunctions")
    await dbfunctions.save_captcha_solved(-100, 11)
    await dbfunctions.save_captcha_solved(-100, 12)
    assert await dbfunctions.has_solved_captcha_once(-100, 11)
    assert await dbfunctions.has_solved_captcha_once(-100, 12)


async def test_captcha_command_updates_admin_setting(greetings):
    reply_text = AsyncMock()
    message = SimpleNamespace(
        chat=SimpleNamespace(id=-100),
        command=["captcha", "开启"],
        text="/captcha 开启",
        reply_text=reply_text,
    )
    await greetings.captcha_state(None, message)
    assert await greetings.is_captcha_on(-100)
    message.command = ["captcha", "关闭"]
    message.text = "/captcha 关闭"
    await greetings.captcha_state(None, message)
    assert not await greetings.is_captcha_on(-100)


async def test_passed_verification_timer_does_not_ban_later_manual_mute(
    greetings, monkeypatch
):
    monkeypatch.setattr(greetings.asyncio, "sleep", AsyncMock())
    chat = SimpleNamespace(
        id=-100,
        get_member=AsyncMock(
            return_value=SimpleNamespace(status=ChatMemberStatus.RESTRICTED)
        ),
        ban_member=AsyncMock(),
    )
    message = SimpleNamespace(chat=chat, delete=AsyncMock())
    user = SimpleNamespace(id=11)
    await greetings.update_captcha_cache(
        [{"chat_id": -100, "user_id": 11, "attempt_id": "old"}]
    )
    await greetings.save_captcha_solved(-100, 11)
    await greetings.kick_restricted_after_delay(300, message, user, "old")

    chat.ban_member.assert_not_awaited()
    message.delete.assert_not_awaited()
    assert await greetings.get_captcha_cache() == []


async def test_success_callback_invalidates_timer_before_manual_mute(
    greetings, monkeypatch
):
    monkeypatch.setattr(greetings.asyncio, "sleep", AsyncMock())
    monkeypatch.setattr(greetings, "send_welcome_message", AsyncMock())
    chat = SimpleNamespace(
        id=-100,
        unban_member=AsyncMock(),
        get_member=AsyncMock(
            return_value=SimpleNamespace(status=ChatMemberStatus.RESTRICTED)
        ),
        ban_member=AsyncMock(),
    )
    button_message = SimpleNamespace(chat=chat, delete=AsyncMock())
    callback = SimpleNamespace(
        data="pressed_button ABCD 11",
        from_user=SimpleNamespace(id=11),
        message=button_message,
        answer=AsyncMock(),
    )
    pending = {
        "chat_id": -100,
        "user_id": 11,
        "attempt_id": "attempt",
        "answer": "ABCD",
        "keyboard": [["answer"]],
        "attempts": 0,
    }
    greetings.answers_dicc = [pending]
    await greetings.update_captcha_cache([pending])

    await greetings.callback_query_welcome_button(None, callback)
    await greetings.kick_restricted_after_delay(
        300, button_message, SimpleNamespace(id=11), "attempt"
    )

    chat.unban_member.assert_awaited_once_with(11)
    chat.ban_member.assert_not_awaited()
    assert await greetings.get_captcha_cache() == []


async def test_timeout_only_bans_matching_pending_verification(greetings, monkeypatch):
    monkeypatch.setattr(greetings.asyncio, "sleep", AsyncMock())
    first_chat = SimpleNamespace(
        id=-100,
        get_member=AsyncMock(
            return_value=SimpleNamespace(status=ChatMemberStatus.RESTRICTED)
        ),
        ban_member=AsyncMock(),
    )
    second_chat = SimpleNamespace(id=-200, ban_member=AsyncMock())
    message = SimpleNamespace(chat=first_chat, delete=AsyncMock())
    await greetings.update_captcha_cache(
        [
            {"chat_id": -100, "user_id": 11, "attempt_id": "new"},
            {"chat_id": -200, "user_id": 11, "attempt_id": "other"},
        ]
    )

    await greetings.kick_restricted_after_delay(
        300, message, SimpleNamespace(id=11), "old"
    )
    first_chat.ban_member.assert_not_awaited()
    await greetings.kick_restricted_after_delay(
        300, message, SimpleNamespace(id=11), "new"
    )
    first_chat.ban_member.assert_awaited_once()
    second_chat.ban_member.assert_not_awaited()
    assert await greetings.get_captcha_cache() == [
        {"chat_id": -200, "user_id": 11, "attempt_id": "other"}
    ]