fix(admin,db): fix /kill and /unban crashes, persist bans, harden add_user

blackjackbot/commands/admin/commands.py

@@ -67,6 +67,7 @@ def kill_game_cmd(update, context):
     """Kills the game for a certain chat/group"""
     if len(context.args) == 0:
         update.message.reply_text("Please provide a chat_id!")
+        return
 
     chat_id = context.args[0]
     # Input validation for chat_id

database/database.py

@@ -112,6 +112,8 @@ def is_user_banned(self, user_id):
 
     def ban_user(self, user_id):
         """Bans a user from using a the bot"""
+        # Make sure the user exists in the database, so that the ban persists across restarts
+        self.cursor.execute("INSERT OR IGNORE INTO users (user_id) VALUES (?);", [str(user_id)])
         self.cursor.execute("UPDATE users SET banned=1 WHERE user_id=?;", [str(user_id)])
         self.connection.commit()
         self._banned_users.add(int(user_id))
@@ -120,7 +122,7 @@ def unban_user(self, user_id):
         """Unbans a user from using a the bot"""
         self.cursor.execute("UPDATE users SET banned=0 WHERE user_id=?;", [str(user_id)])
         self.connection.commit()
-        self._banned_users.remove(int(user_id))
+        self._banned_users.discard(int(user_id))
 
     def get_recent_players(self):
         one_day_in_secs = 60 * 60 * 24
@@ -174,12 +176,16 @@ def add_user(self, user_id, lang_id, first_name, last_name, username):
         self._add_user(user_id, lang_id, first_name, last_name, username)
 
     def _add_user(self, user_id, lang_id, first_name, last_name, username):
+        # Telegram does not always provide a language_code, but the chats table requires one
+        lang_id = lang_id or "en"
         try:
             self.cursor.execute("INSERT INTO users VALUES (?, ?, ?, ?, 0, 0, 0, 0, 0);", [str(user_id), first_name, last_name, username])
-            self.cursor.execute("INSERT INTO chats VALUES (?, ?);", [str(user_id), lang_id])
+            # The chat might already exist, e.g. if the user changed the language via /language before
+            self.cursor.execute("INSERT OR IGNORE INTO chats VALUES (?, ?);", [str(user_id), lang_id])
             self.connection.commit()
         except sqlite3.IntegrityError:
-            return
+            # Don't leave a half written user in the open transaction
+            self.connection.rollback()
 
     def set_games_won(self, games_won, user_id):
         self.cursor.execute("UPDATE users SET games_won = ? WHERE user_id = ?;", [games_won, str(user_id)])

database/tests/database_test.py

@@ -0,0 +1,56 @@
+# -*- coding: utf-8 -*-
+import unittest
+from random import randint
+
+from database import Database
+
+
+class DatabaseTest(unittest.TestCase):
+
+    def setUp(self):
+        self.db = Database()
+        # Use random IDs so that the tests don't interfere with existing data in the database file
+        self.user_id = randint(10 ** 9, 10 ** 10)
+        self.addCleanup(self._delete_user, self.user_id)
+
+    def _delete_user(self, user_id):
+        self.db.cursor.execute("DELETE FROM users WHERE user_id=?;", [str(user_id)])
+        self.db.cursor.execute("DELETE FROM chats WHERE chat_id=?;", [str(user_id)])
+        self.db.connection.commit()
+        self.db.get_banned_users().discard(user_id)
+
+    def test_unban_not_banned_user(self):
+        """Unbanning a user who is not banned must not raise an error"""
+        self.db.unban_user(self.user_id)
+        self.assertFalse(self.db.is_user_banned(self.user_id))
+
+    def test_ban_unknown_user_persists(self):
+        """Banning a user who never used the bot must be stored in the database, not only in memory"""
+        self.db.ban_user(self.user_id)
+
+        self.assertTrue(self.db.is_user_banned(self.user_id))
+        user = self.db.get_user(self.user_id)
+        self.assertIsNotNone(user)
+        self.assertEqual(1, user["banned"])
+
+    def test_add_user_without_language_code(self):
+        """Users without a language_code must be stored with the default language"""
+        self.db.add_user(self.user_id, None, "first", "last", "username")
+
+        self.assertIsNotNone(self.db.get_user(self.user_id))
+        self.db.cursor.execute("SELECT lang_id FROM chats WHERE chat_id=?;", [str(self.user_id)])
+        self.assertEqual("en", self.db.cursor.fetchone()["lang_id"])
+
+    def test_add_user_with_existing_chat(self):
+        """Users whose private chat is already stored (e.g. after using /language) must still be added"""
+        self.db.set_lang_id(chat_id=self.user_id, lang_id="de")
+
+        self.db.add_user(self.user_id, "en", "first", "last", "username")
+
+        self.assertIsNotNone(self.db.get_user(self.user_id))
+        self.db.cursor.execute("SELECT lang_id FROM chats WHERE chat_id=?;", [str(self.user_id)])
+        self.assertEqual("de", self.db.cursor.fetchone()["lang_id"])
+
+
+if __name__ == '__main__':
+    unittest.main()