From 2d4bf992de2166944fa18feab6d457e0d83f15c6 Mon Sep 17 00:00:00 2001
From: Michal Rudowicz <michal.rudowicz@fl9.eu>
Date: Thu, 13 Jun 2024 21:28:13 +0200
Subject: [PATCH] Unit tests for banning

---
 tests.py | 37 ++++++++++++++++++++++++++++++-------
 1 file changed, 30 insertions(+), 7 deletions(-)

diff --git a/tests.py b/tests.py
index 92b0682..379cd07 100644
--- a/tests.py
+++ b/tests.py
@@ -1,7 +1,7 @@
 #!/usr/bin/env python
 
 import unittest
-from unittest.mock import MagicMock
+from unittest.mock import AsyncMock
 import os
 import re
 import sys
@@ -19,20 +19,23 @@ BOT_SPEC = ["kick_chat_member", "delete_message"]
 
 @dataclass
 class Chat:
-    ban_member = MagicMock()
+    ban_member = AsyncMock()
     effective_name = "effective_chat_name"
 
 
-@dataclass
 class User:
-    id = 0
-    name = "User name"
+    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 = MagicMock()
+    delete = AsyncMock()
     chat = Chat()
     from_user = User()
 
@@ -50,17 +53,37 @@ def make_update(msg: str) -> Update:
 
 class TestAutoBanBot(unittest.IsolatedAsyncioTestCase):
     def setUp(self):
-        self.regexes = list(map(lambda r: re.compile(r, re.I), ["test123", "(t\\.me\\/|\\@)[a-z]+bot"]))
+        self.regexes = list(map(lambda r: re.compile(r, re.I), [
+            "test123",
+            "(?:@|(?:(?:(?:https?://)?t(?:elegram)?)\.me\/))(\w{4,})bot"
+        ]))
 
     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(msg=msg):
                 update = make_update(msg)
+                update.message.delete.reset_mock()
+                update.message.chat.ban_member.reset_mock()
                 await autobanbot.new_msg(update, None, self.regexes)
                 update.message.delete.assert_not_called()
                 update.message.chat.ban_member.assert_not_called()
 
+    async def test_not_bannable(self):
+        messages = [
+            "t.me/whatever_bot?not bannable message",
+            "https://www.t.me/ohhaaaaaaaibot/siematest123elo",
+            "@Hello_BoT"
+        ]
+        for msg in messages:
+            with self.subTest(msg=msg):
+                update = make_update(msg)
+                update.message.delete.reset_mock()
+                update.message.chat.ban_member.reset_mock()
+                await autobanbot.new_msg(update, None, self.regexes)
+                update.message.delete.assert_called_once()
+                update.message.chat.ban_member.assert_called_once_with(update.message.from_user.id)
+
 
 if __name__ == '__main__':
     unittest.main()