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"} ]