diff --git a/blackjackbot/commands/admin/commands.py b/blackjackbot/commands/admin/commands.py index 1c8c46f..011afbc 100644 --- a/blackjackbot/commands/admin/commands.py +++ b/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 diff --git a/database/database.py b/database/database.py index d49ab49..82ee9a1 100644 --- a/database/database.py +++ b/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)]) diff --git a/database/tests/database_test.py b/database/tests/database_test.py new file mode 100644 index 0000000..568ca9d --- /dev/null +++ b/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()