From 7c0d5d936c7df1e133f6456cc806702c10145872 Mon Sep 17 00:00:00 2001 From: Eriks K Date: Sun, 19 Apr 2020 13:48:02 +0300 Subject: [PATCH] Improved multi server supprt, some more exception handling, db structure changed --- db.py | 26 ++++++++++++-------------- discord_bot.py | 46 +++++++++++++++++++++++++++------------------- tests.py | 38 +++++++++++++++++++++++--------------- 3 files changed, 62 insertions(+), 48 deletions(-) diff --git a/db.py b/db.py index 280df0f..3e35058 100644 --- a/db.py +++ b/db.py @@ -1,3 +1,4 @@ +from collections import namedtuple from sqlite3 import IntegrityError from typing import List, Union, Dict, Optional, Set @@ -20,12 +21,13 @@ class DiscordDB: self._db.create_table("member", {"id": int, "name": str}, pk="id", not_null={"id", "name"}) else: if "mention_number" in self._db.table('member').columns_dict.keys(): - members: Set[Dict[str, Union[int, str]]] = {*[]} + memb = namedtuple('Member', ['id', 'name']) + members: Set['memb'] = {*[]} for row in self._db.table('member').rows: - members.add({'id': row['mention_number'], 'name': row['name']}) + members.add(memb(row['mention_number'], row['name'])) self._db.table('member').drop() self._db.create_table("member", {"id": int, "name": str}, pk="id", not_null={"id", "name"}) - self._db.table('member').insert_all(list(members)) + self._db.table('member').insert_all([m._asdict() for m in members]) migrate_db = True if "player" not in self._db.table_names(): @@ -113,16 +115,16 @@ class DiscordDB: except NotFoundError: raise NotFoundError("Member with given id not found") - def add_member(self, id: int, name: str) -> int: + def add_member(self, id: int, name: str) -> Dict[str, Union[int, str]]: """Add discord member. :param id: int Discord member ID :param name: Discord member Name """ try: - return self.member.insert({"id": id, "name": name}).last_pk - except IntegrityError: - return id + self.member.insert({"id": id, "name": name}) + finally: + return self.member.get(id) def update_member(self, member_id: int, name: str) -> bool: """Update discord member"s record @@ -136,7 +138,7 @@ class DiscordDB: try: member = self.get_member(member_id) except NotFoundError: - member = self.member.get(self.add_member(member_id, name)) + member = self.add_member(member_id, name) self.member.update(member["id"], {"name": name}) return True @@ -198,9 +200,5 @@ class DiscordDB: def get_hunted_player_ids(self) -> List[int]: return [r["player_id"] for r in self.hunted_players.rows] - def get_members_to_notify(self, pid: int) -> List[int]: - members = [] - for row in self.hunted.rows_where("player_id = ?", [pid]): - member = self.get_member(row["member_id"]) - members.append(member["id"]) - return members + def get_members_to_notify(self, pid: int) -> List[Dict[str, Union[int, str]]]: + return [r for r in self.hunted.rows_where("player_id = ?", [pid])] diff --git a/discord_bot.py b/discord_bot.py index 4036102..e59dce9 100644 --- a/discord_bot.py +++ b/discord_bot.py @@ -26,9 +26,6 @@ os.chdir(os.path.abspath(os.path.dirname(sys.argv[0]))) pid = str(os.getpid()) pidfile = "pid" -if os.path.isfile(pidfile): - print("%s already exists, exiting" % pidfile) - sys.exit() with open(pidfile, 'w') as f: f.write(str(os.getpid())) @@ -89,6 +86,8 @@ def check_player(player_id: int) -> bool: r = requests.get(f'https://www.erepublik.com/en/main/citizen-profile-json/{player_id}').json() except JSONDecodeError: return False + if r.get('error') or not r.get('status'): + return False DB.add_player(player_id, r.get('citizen').get('name')) return True @@ -147,14 +146,14 @@ class MyClient(discord.Client): pid = side_data['citizenId'] medal_key = (pid, battle['id'], div['div'], battle[side]['id'], side_data['damage']) if not DB.check_medal(*medal_key): - for member in DB.get_members_to_notify(pid): - format_data = dict(author=member, player=DB.get_player(pid)['name'], + for hunt_row in DB.get_members_to_notify(pid): + format_data = dict(author=hunt_row['member_id'], player=DB.get_player(pid)['name'], battle=bid, region=battle.get('region').get('name'), division=div['div'], dmg=side_data['damage'], side=COUNTRIES[battle[side]['id']]) - await self.get_channel(603527159109124096).send( + await self.get_channel(hunt_row['channel_id']).send( "<@{author}> **{player}** detected in battle for {region} on {side} side in d{division} with {dmg:,d}dmg\n" "https://www.erepublik.com/en/military/battlefield/{battle}".format( **format_data) @@ -252,19 +251,21 @@ if __name__ == "__main__": try: local_member_id = DB.get_member(ctx.author.id).get('id') except NotFoundError: - local_member_id = DB.add_member(ctx.author.id, ctx.author.name) - if DB.add_hunted_player(player_id, local_member_id): - await ctx.send(f"{ctx.author.mention} You'll be notified for all **{player_name}** medals") + local_member_id = DB.add_member(ctx.author.id, ctx.author.name).get('id') + if ctx.channel.type.value == 1: + await ctx.send(f"{ctx.author.mention}, sorry, but currently I'm unable to notify You in DM channel!") + elif DB.add_hunted_player(player_id, local_member_id, ctx.channel.id): + await ctx.send(f"{ctx.author.mention} You'll be notified for **{player_name}** medals in this channel") else: - await ctx.send(f"{ctx.author.mention} You are already being notified for all **{player_name}** medals") + await ctx.send(f"{ctx.author.mention} You are already being notified for **{player_name}** medals") @bot.command(description="Informēt par spēlētāja mēģinājumiem ņemt medaļas", help="Piereģistrēties uz spēlētāja medaļu paziņošanu", category="Hunting") async def my_hunt(ctx): msgs = [] - for hunt in DB.get_member_hunted_players(ctx.author.id): - msgs.append(f"`{hunt['id']}` - **{hunt['name']}**") + for hunted_player in DB.get_member_hunted_players(ctx.author.id): + msgs.append(f"`{hunted_player['id']}` - **{hunted_player['name']}**") if msgs: msg = "\n".join(msgs) await ctx.send(f"{ctx.author.mention} You are hunting:\n{msg}") @@ -279,16 +280,23 @@ if __name__ == "__main__": await ctx.send(f"{ctx.author.mention} didn't find any player with `id: {player_id}`!") else: player_name = DB.get_player(player_id).get('name') - local_member_id = DB.get_member(ctx.author.id).get('id') + try: + local_member_id = DB.get_member(ctx.author.id).get('id') + except NotFoundError: + local_member_id = DB.add_member(ctx.author.id, ctx.author.name).get('id') if DB.remove_hunted_player(player_id, local_member_id): await ctx.send(f"{ctx.author.mention} You won't be notified for **{player_name}** medals") else: await ctx.send(f"{ctx.author.mention} You were not hunting **{player_name}** medals") - try: - loop.create_task(bot.start(DISCORD_TOKEN)) - loop.create_task(client.start(DISCORD_TOKEN)) - loop.run_forever() - finally: - os.unlink(pidfile) + + @hunt.error + @remove_hunt.error + async def hunt_error(ctx, error): + if isinstance(error, commands.BadArgument): + await ctx.send('spēlētāja identifikators jāpadod kā skaitliska vērtība, piemēram, 1620414') + + loop.create_task(bot.start(DISCORD_TOKEN)) + loop.create_task(client.start(DISCORD_TOKEN)) + loop.run_forever() diff --git a/tests.py b/tests.py index 9a15816..fa2acf3 100644 --- a/tests.py +++ b/tests.py @@ -12,7 +12,7 @@ class TestDatabase(unittest.TestCase): def test_member(self): member = {'id': 1200, 'name': 'username'} self.db.add_member(**member) - self.assertEqual(self.db.add_member(**member), member['id']) + self.assertEqual(self.db.add_member(**member), member) self.assertRaises(NotFoundError, self.db.get_member, member_id=100) self.assertEqual(self.db.get_member(member_id=member['id']), member) @@ -41,20 +41,28 @@ class TestDatabase(unittest.TestCase): self.assertTrue(self.db.delete_medals([kwargs['bid']])) def test_hunt(self): - member_id = self.db.add_member(2, name="one") - member_id2 = self.db.add_member(3, name="two") - member_id3 = self.db.add_member(4, name="three") + member_1 = self.db.add_member(2, name="one") + member_2 = self.db.add_member(3, name="two") + member_3 = self.db.add_member(4, name="three") self.db.add_player(1, 'plato') self.db.add_player(2, 'draco') - self.assertFalse(self.db.check_hunt(1, member_id)) - self.assertFalse(self.db.remove_hunted_player(1, member_id)) - self.assertTrue(self.db.add_hunted_player(1, member_id, 123)) - self.assertTrue(self.db.add_hunted_player(1, member_id2, 234)) - self.assertTrue(self.db.add_hunted_player(1, member_id3, 345)) - self.assertTrue(self.db.add_hunted_player(2, member_id, 456)) + self.assertFalse(self.db.check_hunt(1, member_1['id'])) + self.assertFalse(self.db.remove_hunted_player(1, member_1['id'])) + player_1_hunt = [{'id': 1, "member_id": member_1['id'], 'player_id': 1, 'channel_id': 123}, + {'id': 2, "member_id": member_2['id'], 'player_id': 1, 'channel_id': 234}, + {'id': 3, "member_id": member_3['id'], 'player_id': 1, 'channel_id': 345}] + for hunt in player_1_hunt: + self.assertTrue(self.db.add_hunted_player(hunt['player_id'], hunt['member_id'], hunt['channel_id'])) + + player_2_hunt = [{'id': 4, "member_id": member_1['id'], 'player_id': 2, 'channel_id': 456}] + for hunt in player_2_hunt: + self.assertTrue(self.db.add_hunted_player(hunt['player_id'], hunt['member_id'], hunt['channel_id'])) + self.assertListEqual(self.db.get_hunted_player_ids(), [1, 2]) - self.assertListEqual(self.db.get_members_to_notify(1), [2, 3, 4]) - self.assertListEqual(self.db.get_members_to_notify(2), [2]) - self.assertFalse(self.db.add_hunted_player(1, member_id, 567)) - self.assertTrue(self.db.check_hunt(1, member_id)) - self.assertTrue(self.db.remove_hunted_player(1, member_id)) \ No newline at end of file + self.assertListEqual(self.db.get_members_to_notify(1), player_1_hunt) + self.assertListEqual(self.db.get_members_to_notify(2), player_2_hunt) + + self.assertFalse(self.db.add_hunted_player(1, member_1['id'], 567)) + self.assertTrue(self.db.check_hunt(1, member_1['id'])) + self.assertTrue(self.db.remove_hunted_player(1, member_1['id'])) + self.assertFalse(self.db.check_hunt(1, member_1['id']))