תבנית לכתיבת בדיקה לבוט ב python-telegram-bot

16/07/2023

ספריית python-telegram-bot מציעה ממשק מאוד מקיף לבניית בוטים לטלגרם. הספריה מגיעה עם אינסוף דוגמאות ותיעוד מפורט, ולמרות זאת הדף על כתיבת בדיקות בויקי שלהם כמעט ריק. אחרי קצת חפירה בקוד שלהם הגעתי לתבנית הבאה לבדיקת יחידה ב pytest לבוט שכתוב ב python-telegram-bot. נתחיל בקוד ואחריו ההסברים:

from telegram import Message, MessageEntity, User, Chat, Update
import pytest
from unittest.mock import MagicMock
from datetime import datetime
import client.telegram.bot as bot

class AsyncMock(MagicMock):
    async def __call__(self, *args, **kwargs):
        return super(AsyncMock, self).__call__(*args, **kwargs)


@pytest.mark.asyncio
async def test_create_user(monkeypatch):
    with monkeypatch.context() as m:
        app = bot.create_application()
        await app.initialize()

        message = Message(
            message_id=0,
            date=datetime.utcnow(),
            chat=Chat(3, "private"),
            from_user=User(first_name='Misses Test', id=123, is_bot=False),
            text="/start",
            entities=[MessageEntity(type=MessageEntity.BOT_COMMAND,
                                    offset=0, length=len('/start'))]
        )
        message.set_bot(app.bot)

        send_message = AsyncMock()
        m.setattr(app.bot, "_send_message", send_message)

        await app.process_update(Update(update_id=0, message=message))
        assert 'Hello World' in send_message.call_args[0][1]['text']

והקוד עבור הבוט המתאים, מתוך קובץ שבדוגמה שלנו נקרא client/telegram/bot.py יהיה:

async def start_conversation(update: Update, context: ContextTypes.DEFAULT_TYPE):
    await update.message.reply_text(f"Hello World")


def create_application():
    # Create the Application and pass it your bot's token.
    token = os.environ.get("TELEGRAM_BOT_TOKEN")
    application = Application.builder().token(token).build()

    application.add_handler(CommandHandler('start', start_conversation))
    return application


def main() -> None:
    """Start the bot."""
    application = create_application()
    # Run the bot until the user presses Ctrl-C
    application.run_polling()


if __name__ == "__main__":
    print("Starting...")
    main()

בואו נראה מה היה לנו:

  1. קוד הבדיקה הוא פונקציה אסינכרונית. בשביל שזה יעבוד יש להתקין את הפלאגין pytest-asyncio.

  2. הבדיקה צריכה לזייף הודעת /start ואז צריכה לוודא שהבוט ענה בהודעת Hello World. בשביל זה אנחנו ממש יוצרים את הבוט (תחילת קוד הבדיקה) ומפעילים את app.initialize שזו פונקציה של python-telegram-bot שדואגת לכל מיני איתחולים. בחיבור רגיל לטלגרם קריאה ל run_polling מפעילה את ה initialize אוטומטית.

  3. הפונקציה process_update היא נקודת הכניסה לבוט והיא מקבלת אוביקט הודעה ומנתבת אותו לפונקציית הטיפול המתאימה לפי הכללים שהגדיר הבוט (אלה פקודות ה add_handler שרואים בקוד הבוט). בקוד הבדיקה משתמשים בה כדי לשלוח את הודעת ה /start.

  4. יצירת אוביקט ההודעה בקוד לוקחת את רוב שורות הקוד בבדיקה. כן צריך להוציא את זה ל fixture עם ערכי ברירת מחדל טובים וזה יכול לחסוך כמה תווים בכל בדיקה. שני הדברים החשובים כאן הם טקסט ההודעה ורשימת הישויות שבה, כי אלה הדברים שהבוט שלנו יראה.

  5. הפקודה set_bot היא שמאפשרת בתוך קוד הבוט לגשת ל update.message.reply_text, ולכל שאר הפונקציות של הבוט דרך אוביקט ההודעה.

  6. בשביל לקבל בחזרה לבדיקה את הפרמטרים שנשלחו לפונקציות שליחת ההודעה (אלה שהיו אמורים להישלח לטלגרם בהרצה אמיתית של הבוט), אני יוצר MagicMock ושומר אותו בשדה _send_message של אוביקט הבוט. הפונקציה _send_message נועדה לשימוש פנימי ולכן אפשר לדרוס אותה. ניסיון לדרוס פונקציות ציבוריות של הבוט ייכשל בגלל הקוד הבא בתוך ספריית python-telegram-bot באוביקט הבוט:

def __setattr__(self, key: str, value: object) -> None:
    """Overrides :meth:`object.__setattr__` to prevent the overriding of attributes.

    Raises:
        :exc:`AttributeError`
    """
    # protected attributes can always be set for convenient internal use
    if key[0] == "_" or not getattr(self, "_frozen", True):
        super().__setattr__(key, value)
        return

    raise AttributeError(
        f"Attribute `{key}` of class `{self.__class__.__name__}` can't be set!"
    )
  1. אחרי דריסת _send_message נשאר רק לשלוח את ההודעה ל process_update ולוודא את ההודעה שהבוט שולח דרך בדיקה של הקריאות שבוצעו ל Mock בשורה האחרונה של הבדיקה.