"""Access-control tests. Run with: python test_access.py Covers who gets into the bot in public mode, and the membership cache that keeps inline search from hammering the Telegram API on every keystroke. """ import asyncio import importlib import logging import os import sys import time os.environ.setdefault("BOT_TOKEN", "123456789:AAHtestTokenPlaceholder_ABCDEFGHIJKLMNOP") os.environ.setdefault("DB_PATH", ":memory:") sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from aiogram.exceptions import TelegramAPIError # Several cases deliberately make the API fail; the bot logs those with a # traceback, which is correct in production and only noise here. logging.disable(logging.CRITICAL) CONFIG_KEYS = ("PUBLIC_MODE", "REQUIRED_CHATS", "REQUIRE_MODE", "ALLOWED_USERS", "ADMIN_USERS", "ACCEPT_SUBMISSIONS") class FakeMember: def __init__(self, status: str, is_member: bool | None = None): self.status = status if is_member is not None: self.is_member = is_member def load_bot(**env): """Reload the bot module with a fresh environment; config is read at import.""" for key in CONFIG_KEYS: os.environ.pop(key, None) os.environ.update(env) import bot return importlib.reload(bot) def stub_membership(B, members: dict, calls: list | None = None): async def get_chat_member(chat, user_id): if calls is not None: calls.append((chat, user_id)) if (chat, user_id) not in members: raise TelegramAPIError(method=None, message="user not found") return members[(chat, user_id)] B.bot.get_chat_member = get_chat_member async def test_member_statuses(): B = load_bot(PUBLIC_MODE="true", REQUIRED_CHATS="@chan") allowed = ["creator", "administrator", "member"] denied = ["left", "kicked"] for status in allowed: B._membership_cache.clear() stub_membership(B, {("@chan", 1): FakeMember(status)}) assert await B.has_access(1) is True, f"{status} should be allowed" for status in denied: B._membership_cache.clear() stub_membership(B, {("@chan", 1): FakeMember(status)}) assert await B.has_access(1) is False, f"{status} should be denied" # A restricted user is still a member unless they have actually left. B._membership_cache.clear() stub_membership(B, {("@chan", 1): FakeMember("restricted", is_member=True)}) assert await B.has_access(1) is True, "restricted but still in the chat" B._membership_cache.clear() stub_membership(B, {("@chan", 1): FakeMember("restricted", is_member=False)}) assert await B.has_access(1) is False, "restricted and gone" print(" member statuses are classified correctly, including restricted") async def test_require_all_and_any(): members = {("@a", 1): FakeMember("member")} # in @a, not in @b B = load_bot(PUBLIC_MODE="true", REQUIRED_CHATS="@a,@b", REQUIRE_MODE="all") stub_membership(B, members) assert await B.has_access(1) is False, "all mode needs every chat" B = load_bot(PUBLIC_MODE="true", REQUIRED_CHATS="@a,@b", REQUIRE_MODE="any") stub_membership(B, members) assert await B.has_access(1) is True, "any mode needs just one" print(" REQUIRE_MODE all and any behave differently as configured") async def test_cache_absorbs_repeat_checks(): B = load_bot(PUBLIC_MODE="true", REQUIRED_CHATS="@chan") calls: list = [] stub_membership(B, {("@chan", 1): FakeMember("member")}, calls) for _ in range(10): assert await B.has_access(1) is True assert len(calls) == 1, f"10 inline keystrokes caused {len(calls)} API calls" print(" repeated checks hit the API once, not once per keystroke") async def test_denials_expire_sooner_than_grants(): B = load_bot(PUBLIC_MODE="true", REQUIRED_CHATS="@chan") assert B.MEMBERSHIP_TTL_DENIED < B.MEMBERSHIP_TTL_ALLOWED stub_membership(B, {("@chan", 1): FakeMember("member")}) await B.has_access(1) await B.has_access(2) # unknown user -> denied now = time.monotonic() granted = B._membership_cache[(1, "@chan")][1] - now denied = B._membership_cache[(2, "@chan")][1] - now assert denied < granted, "a user who just joined should not wait the full window" print(" denials expire quickly so joining takes effect promptly") async def test_admin_bypasses_membership(): B = load_bot(PUBLIC_MODE="true", REQUIRED_CHATS="@chan", ADMIN_USERS="111") calls: list = [] stub_membership(B, {}, calls) # nobody is a member assert await B.has_access(111) is True, "admins must not be locked out" assert calls == [], "admin check should not call the API at all" assert await B.has_access(222) is False print(" admins bypass the membership check without an API call") async def test_private_mode_still_uses_the_allowlist(): B = load_bot(ALLOWED_USERS="111,222") calls: list = [] stub_membership(B, {}, calls) assert await B.has_access(111) is True assert await B.has_access(333) is False assert calls == [], "a private bot should never call getChatMember" print(" with PUBLIC_MODE off, the allowlist behaves exactly as before") async def test_api_failure_fails_closed(): B = load_bot(PUBLIC_MODE="true", REQUIRED_CHATS="@chan") async def always_fails(chat, user_id): raise TelegramAPIError(method=None, message="bot is not an admin there") B.bot.get_chat_member = always_fails assert await B.has_access(1) is False, "an unreadable chat must not open the gate" print(" a failing membership check denies access rather than opening up") async def test_open_public_bot(): B = load_bot(PUBLIC_MODE="true") # no REQUIRED_CHATS assert await B.has_access(999) is True print(" public mode with no required chats is open to everyone") async def main(): for test in (test_member_statuses, test_require_all_and_any, test_cache_absorbs_repeat_checks, test_denials_expire_sooner_than_grants, test_admin_bypasses_membership, test_private_mode_still_uses_the_allowlist, test_api_failure_fails_closed, test_open_public_bot): await test() print("\nAll access tests passed.") if __name__ == "__main__": asyncio.run(main())