memes_bot/test_vault.py

"""Vault tests. Run with: python test_vault.py

The vault stores pointers to messages the bot already holds, so the things
worth pinning down are that albums are grouped into one post, that retrieval
asks Telegram for the right message ids, and that one user can never reach
another user's saved posts.
"""
import asyncio
import datetime
import importlib
import logging
import os
import sys
import tempfile

os.environ.setdefault("BOT_TOKEN", "123456789:AAHtestTokenPlaceholder_ABCDEFGHIJKLMNOP")
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from aiogram.types import Chat, Message, PhotoSize, User

logging.disable(logging.CRITICAL)

CHAT = Chat(id=500, type="private")
USER = User(id=111, is_bot=False, first_name="Owner")
NOW = datetime.datetime.now()


def load(path: str, **env):
    for key in ("PUBLIC_MODE", "REQUIRED_CHATS", "ALLOWED_USERS", "ADMIN_USERS"):
        os.environ.pop(key, None)
    os.environ["DB_PATH"] = path
    os.environ.update(env)
    import db, bot
    importlib.reload(db)
    B = importlib.reload(bot)
    B.db = db
    return db, B


def photo_msg(message_id: int, caption: str | None = None, group: str | None = None) -> Message:
    photo = PhotoSize(file_id=f"F{message_id}", file_unique_id=f"U{message_id}", width=1, height=1)
    return Message(message_id=message_id, date=NOW, chat=CHAT, from_user=USER,
                   photo=[photo], caption=caption, media_group_id=group)


# ── database ─────────────────────────────────────────────────────────────────

async def test_roundtrip(db):
    item_id = await db.add_vault_item(111, 500, [12, 10, 11], "A trip to the sea", ["sea"])
    item = await db.get_vault_item(111, item_id)
    assert db.unpack_message_ids(item["message_ids"]) == [12, 10, 11]
    assert item["caption"] == "A trip to the sea"
    print("  a saved post round-trips with its message ids and caption")


async def test_search_by_keyword_and_caption(db):
    await db.add_vault_item(111, 500, [20], "Recipe for lasagne", ["food"])
    assert len(await db.search_vault(111, "food")) == 1, "keyword search"
    assert len(await db.search_vault(111, "lasagne")) == 1, "caption search"
    assert len(await db.search_vault(111, "lasagne")) == 1
    assert await db.search_vault(111, "nonsense") == []
    print("  posts are findable by keyword and by words in the original caption")


async def test_vault_is_private_to_its_owner(db):
    mine = await db.add_vault_item(111, 500, [30], "private note", ["secret"])

    assert await db.get_vault_item(222, mine) is None, "another user fetched my post"
    assert await db.search_vault(222, "secret") == [], "another user searched my vault"
    assert await db.delete_vault_item(222, mine) is False, "another user deleted my post"
    assert await db.get_vault_item(111, mine) is not None, "owner lost their own post"
    print("  one user cannot read, search or delete another user's vault")


async def test_delete(db):
    item_id = await db.add_vault_item(111, 500, [40], "", ["bye"])
    assert await db.delete_vault_item(111, item_id) is True
    assert await db.get_vault_item(111, item_id) is None
    assert await db.delete_vault_item(111, item_id) is False
    print("  deleting removes the post and is safe to repeat")


# ── album grouping ───────────────────────────────────────────────────────────

async def test_album_becomes_one_post(B):
    B.ALBUM_DEBOUNCE = 0.05
    B._album_buffers.clear()
    B._pending_forwards.clear()
    sent: list = []

    async def send_message(chat_id, text, **kw):
        sent.append(text)
    B.bot.send_message = send_message

    # Telegram delivers an album as separate messages; only one has the caption.
    for msg in (photo_msg(10, caption="beach day", group="g1"),
                photo_msg(11, group="g1"),
                photo_msg(12, group="g1")):
        await B.on_forwarded(msg, None)
    await asyncio.sleep(0.3)

    assert len(sent) == 1, f"an album should prompt once, got {len(sent)} prompts"
    assert len(B._pending_forwards) == 1
    record = next(iter(B._pending_forwards.values()))
    assert record["message_ids"] == [10, 11, 12]
    assert record["caption"] == "beach day", "the one caption in the album was lost"
    assert len(record["media"]) == 3, "every item in the album should be extractable"
    print("  an album is buffered into one post, keeping its caption and all media")


async def test_single_forward_prompts_immediately(B):
    B._album_buffers.clear()
    B._pending_forwards.clear()
    sent: list = []

    async def send_message(chat_id, text, **kw):
        sent.append(text)
    B.bot.send_message = send_message

    await B.on_forwarded(photo_msg(20, caption="just one"), None)
    assert len(sent) == 1, "a lone forward should not wait for an album"
    record = next(iter(B._pending_forwards.values()))
    assert record["message_ids"] == [20]
    print("  a single forwarded post prompts straight away")


async def test_pending_prompts_are_bounded(B):
    B._pending_forwards.clear()
    B.bot.send_message = lambda *a, **kw: asyncio.sleep(0)

    for i in range(B.PENDING_FORWARD_LIMIT + 25):
        await B.prompt_forward_choice(
            {"user_id": 1, "chat_id": 500, "message_ids": [i], "caption": "", "media": []}
        )
    assert len(B._pending_forwards) <= B.PENDING_FORWARD_LIMIT
    print("  abandoned prompts cannot grow without bound")


async def test_retrieval_forwards_ascending_ids(db, B):
    item_id = await db.add_vault_item(111, 500, [12, 10, 11], "out of order", ["x"])
    forwarded: dict = {}

    async def forward_messages(chat_id, from_chat_id, message_ids):
        forwarded.update(chat_id=chat_id, from_chat_id=from_chat_id, message_ids=message_ids)
    B.bot.forward_messages = forward_messages

    class Cb:
        data = f"vault:get:{item_id}"
        from_user = type("U", (), {"id": 111})()
        message = type("M", (), {"chat": CHAT})()
        async def answer(self, *a, **kw): pass

    await B.on_vault_action(Cb())
    # Telegram requires strictly increasing ids, and the album must stay whole.
    assert forwarded["message_ids"] == [10, 11, 12], forwarded
    assert forwarded["from_chat_id"] == 500
    print("  retrieval forwards the whole post with ids in ascending order")


async def main():
    with tempfile.TemporaryDirectory() as tmp:
        db, B = load(os.path.join(tmp, "vault.db"), ADMIN_USERS="111")
        await db.init_db()

        for test in (test_roundtrip, test_search_by_keyword_and_caption,
                     test_vault_is_private_to_its_owner, test_delete):
            await test(db)
        for test in (test_album_becomes_one_post, test_single_forward_prompts_immediately,
                     test_pending_prompts_are_bounded):
            await test(B)
        await test_retrieval_forwards_ascending_ids(db, B)
    print("\nAll vault tests passed.")


if __name__ == "__main__":
    asyncio.run(main())