auto-ban-bot/tests.py

#!/usr/bin/env python

import unittest
from unittest.mock import AsyncMock
import os
import re
import sys
from dataclasses import dataclass
fakes_dir = os.path.join(os.path.dirname(__file__), 'fakes')
assert(os.path.exists(fakes_dir))
sys.path.insert(0, fakes_dir)
import autobanbot


ALLOWED_CHAT_ID = 123
DISALLOWED_CHAT_ID = 789


@dataclass
class Chat:
    id: int = 0
    ban_member = AsyncMock()
    effective_name = "effective_chat_name"


class User:
    id_count = 0

    def __init__(self):
        self.id = User.id_count
        self.name = "User name"
        User.id_count += 1


@dataclass
class Message:
    text = "Message text"
    delete = AsyncMock()
    chat = Chat()
    from_user = User()


@dataclass
class Update:
    message = Message()
    edited_message = None


def make_update(msg: str, chat_id: int) -> Update:
    update = Update()
    update.message.text = msg
    update.message.chat.id = chat_id
    return update

def make_edited_message(msg: str, chat_id: int) -> Update:
    update = Update()
    update.message = None
    update.edited_message = Message()
    update.edited_message.text = msg
    update.edited_message.chat.id = chat_id
    return update


class TestAutoBanBot(unittest.IsolatedAsyncioTestCase):
    def setUp(self):
        config = autobanbot.load_config("config.json.example")
        self.regexes = config['regexes']
        self.allowed_chats = config['allowedChats']

    async def test_not_bannable(self):
        messages = ["not bannable message", "siematest_123elo", "i am not a bot", "t.me, i'm not a bot"]
        for msg in messages:
            with self.subTest("new message", message=msg):
                update = make_update(msg, ALLOWED_CHAT_ID)
                self.assertIsNone(update.edited_message)
                update.message.delete.reset_mock()
                update.message.chat.ban_member.reset_mock()
                await autobanbot.new_msg(update, None, self.regexes, self.allowed_chats)
                update.message.delete.assert_not_called()
                update.message.chat.ban_member.assert_not_called()
            with self.subTest("message edit", message=msg):
                update = make_edited_message(msg, ALLOWED_CHAT_ID)
                self.assertIsNone(update.message)
                update.edited_message.delete.reset_mock()
                update.edited_message.chat.ban_member.reset_mock()
                await autobanbot.new_msg(update, None, self.regexes, self.allowed_chats)
                update.edited_message.delete.assert_not_called()
                update.edited_message.chat.ban_member.assert_not_called()

    async def test_bannable(self):
        messages = [
            "hitest123hello",
            "HiTeSt123HeLlO",
            "t.me/whatever_bot?not bannable message",
            "https://www.t.me/ohhaaaaaaaibot/siematest123elo",
            "@Hello_BoT"
        ]
        for msg in messages:
            with self.subTest("Allowed chat causes a ban", message=msg):
                update = make_update(msg, ALLOWED_CHAT_ID)
                self.assertIsNone(update.edited_message)
                update.message.delete.reset_mock()
                update.message.chat.ban_member.reset_mock()
                await autobanbot.new_msg(update, None, self.regexes, self.allowed_chats)
                update.message.delete.assert_called_once()
                update.message.chat.ban_member.assert_called_once_with(update.message.from_user.id)
            with self.subTest("Unknown chat is ignored", message=msg):
                update = make_update(msg, DISALLOWED_CHAT_ID)
                self.assertIsNone(update.edited_message)
                update.message.delete.reset_mock()
                update.message.chat.ban_member.reset_mock()
                await autobanbot.new_msg(update, None, self.regexes, self.allowed_chats)
                update.message.delete.assert_not_called()
                update.message.chat.ban_member.assert_not_called()
            with self.subTest("Edit in allowed chat causes a ban", message=msg):
                update = make_edited_message(msg, ALLOWED_CHAT_ID)
                self.assertIsNone(update.message)
                update.edited_message.delete.reset_mock()
                update.edited_message.chat.ban_member.reset_mock()
                await autobanbot.new_msg(update, None, self.regexes, self.allowed_chats)
                update.edited_message.delete.assert_called_once()
                update.edited_message.chat.ban_member.assert_called_once_with(update.edited_message.from_user.id)
            with self.subTest("Edit in unknown chat is ignored", message=msg):
                update = make_edited_message(msg, DISALLOWED_CHAT_ID)
                self.assertIsNone(update.message)
                update.edited_message.delete.reset_mock()
                update.edited_message.chat.ban_member.reset_mock()
                await autobanbot.new_msg(update, None, self.regexes, self.allowed_chats)
                update.edited_message.delete.assert_not_called()
                update.edited_message.chat.ban_member.assert_not_called()


if __name__ == '__main__':
    unittest.main()