#!/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()