memes_bot/test_access.py

"""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())