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