class TestKfpDb(): def setup_method(self, method): self.guild_id = 123 self.database = KfpDb(dbFile=":memory:") self.database.add_member(default_user_id) # add a default member def teardown_method(self, method): self.database.teardown() for file in os.listdir(): if file.endswith("test.db"): os.remove(file) def test_addMember(self): self.database.add_member(12346) member = self.database.get_member(12346) assert member.member_id == 12346 def test_dataIntegrity(self): self.database = KfpDb(dbFile="tmp_test.db") self.database.add_member(12346) self.database = KfpDb(dbFile="tmp_test.db") member = self.database.get_member(12346) assert member.member_id == 12346 def test_getMember_notExist(self): assert not self.database.get_member(100) # user 100 does not exist def test_addMultipleMembers(self): member_ids = [1, 2, 3, 4, 5, 6, 7, 8] self.database.add_members(member_ids) for member_id in member_ids: member = Member.get_by_id(member_id) assert member.member_id == member_id def test_memberHasUniqueId(self): with pytest.raises(peewee.IntegrityError): self.database.add_member(default_user_id) def test_increaseExp_notExist(self): assert -1 == self.database.increase_exp(0, 0, 100, 10) # user 100 does not exist def test_increaseExp(self): self.database.increase_exp(0, 0, default_user_id, 10) member = Member.get_by_id(default_user_id) assert member.exp == 10 def test_ignore_increaseExp(self): self.database.set_ignore_xp_channel(3, 5) self.database.increase_exp(3, 5, default_user_id, 10) member = Member.get_by_id(default_user_id) assert member.exp == 0 def test_after_remove_ignore_increaseExp(self): self.database.remove_ignore_xp_channel(3, 5) self.database.set_ignore_xp_channel(3, 5) self.database.remove_ignore_xp_channel(3, 5) self.database.increase_exp(3, 5, default_user_id, 10) member = Member.get_by_id(default_user_id) assert member.exp == 10 def test_rankUp(self): member = Member.get_by_id(default_user_id) assert member.rank == 0 self.database.increase_exp(0, 0, default_user_id, 100) member = Member.get_by_id(default_user_id) assert member.rank == 1 def test_addCoin_notExist(self): assert not self.database.add_coin(100, 10) # user 100 does not exist def test_addCoin(self): self.database.add_coin(default_user_id, 10) member = self.database.get_member(default_user_id) assert member.coin == 10 def test_subtractCoin(self): self.database.add_coin(default_user_id, 100) self.database.add_coin(default_user_id, -10) member = self.database.get_member(default_user_id) assert member.coin == 90 def test_notEnoughMoney(self): self.database.add_coin(default_user_id, 100) assert not self.database.add_coin(default_user_id, -1000) member = self.database.get_member(default_user_id) assert member.coin == 100 def test_getMemberRankOrder_number1(self): # 只有一人的時候那人是第一名 assert self.database.get_member_rank_order(0) == 1 def test_getMemberRankOrder_last(self): self.database.add_member(1) self.database.add_member(2) self.database.increase_exp(0, 0, 1, 100) self.database.increase_exp(0, 0, 2, 100) assert self.database.get_member_rank_order(0) == 3 def test_getMemberRankOrder_first(self): self.database.add_member(1) self.database.add_member(2) self.database.increase_exp(0, 0, default_user_id, 101) self.database.increase_exp(0, 0, 1, 100) self.database.increase_exp(0, 0, 2, 100) assert self.database.get_member_rank_order(default_user_id) == 1 def test_sameMemberRankOrder_second(self): self.database.add_member(1) self.database.add_member(2) self.database.add_member(3) self.database.increase_exp(0, 0, default_user_id, 100) self.database.increase_exp(0, 0, 1, 100) self.database.increase_exp(0, 0, 2, 1000) self.database.increase_exp(0, 0, 3, 10) # 經驗相同時排名相同, 並列第二 assert self.database.get_member_rank_order(default_user_id) == 2 assert self.database.get_member_rank_order(1) == 2 def test_setRankupChannel_notExist(self): assert not self.database.get_rankup_channel_id(11) def test_setRankupChannel(self): self.database.set_rankup_channel(11, 123) assert self.database.get_rankup_channel_id(11) == 123 self.database.set_rankup_channel(11, 456) assert self.database.get_rankup_channel_id(11) == 456 def test_resetEveryoneToken(self): self.database.add_member(1) self.database.add_member(2) self.database.add_member(3) self.database.reset_everyone_token() assert self.database.get_member(1).token == 100 assert self.database.get_member(2).token == 100 assert self.database.get_member(3).token == 100
class NewProfile(commands.Cog): #TODO: add check permiision function, base on roles db = None __channels = [] def __init__(self, client, dbFile: str, isTest=False): self.bot = client self.db = KfpDb(dbFile) self.isTest = isTest @commands.Cog.listener('on_message') async def profile_on_message(self, message: Message): if message.author.bot: return if message.channel == None or not message.channel.guild.id in whitelist or message.author.bot: return if self.populateChannels(message, self.isTest): return if not self.channelAllowed(message.channel.id, self.isTest): return member: Member = self.db.get_member(message.author.id) if not member: self.db.add_member(message.author.id) member = self.db.get_member(message.author.id) increaseNumber = randint(10, 25) rank = self.db.increase_exp(message.channel.guild.id, message.channel.id, message.author.id, increaseNumber) assert rank > 0, 'method increase_xp should not retrun less than 1 in profile_on_message' if member.rank != rank: channel = ChannelUtil.getMessageChannelId(message.guild.id) if channel == None: channelToUse = message.channel else: channelToUse = message.guild.get_channel(channel) RPGCharacterUtil.levelUpCharacter(message.author.id, member.rank, rank) await channelToUse.send('恭喜<@{}> 等級提升至{}。'.format( message.author.id, rank)) await self.updateUserKfpRoles(message, rank, channelToUse) self.db.increase_coin(message.guild.id, message.author.id, increaseNumber) @commands.group(name='profile', invoke_without_command=True) async def profile_group(self, ctx: commands.Context, *attr): if not isWhiteList(ctx): if ctx.guild: print( "{} is not on white list, if you are a developer, add your server to the white list" .format(ctx.guild.id)) return if not ChannelUtil.hasChannel(ctx.guild.id, ctx.channel.id, Util.ChannelType.PROFILE): print( "WARNING: cannot run on this channel, if you are a developer, try to use '!commandControll add profile'" ) return memberRow: Member = self.db.get_member(ctx.author.id) if memberRow == None: self.db.add_member(ctx.author.id) memberRow = self.db.get_member(ctx.author.id) profileByte = None iconData = None bgData = None avatar_url = ctx.author.avatar_url_as(format='jpg', size=1024) if avatar_url._url != None: iconData = await avatar_url.read() banner_url = ctx.guild.banner_url if banner_url._url != None: bgData = await banner_url.read() with ProfileImage() as pf: if iconData: pf.setIcon(iconData) if bgData: pf.setBackGround(bgData) pf.setCoin(memberRow.coin) pf.setXp(memberRow.exp) pf.setLevelNumber(memberRow.rank) pf.setRankNumber(self.db.get_member_rank_order(ctx.author.id)) pf.setMemberName(ctx.author.display_name, ctx.author.name) profileByte = pf.generateProfileImage() discordFile = discord.File(io.BytesIO(profileByte), filename='profile.png') await ctx.channel.send(file=discordFile) @profile_group.command(name='bind') @commands.check(isWhiteList) async def profile_group_bind_command(self, ctx: commands.Context, *arg): channel = ctx.channel ChannelUtil.setRankupChannel(ctx.guild.id, channel.id) await channel.send('<@!{}> 設定升級訊息將會於此。'.format(ctx.author.id)) @profile_group.command(name='allowed') async def profile_allowed_channels_command(self, ctx: commands.Context, *arg): msg = "```" msg += "allowed channel list:\n" for channel_id in self.__channels: channel = ctx.guild.get_channel(channel_id) msg += f"{channel.id}: {channel.name}\n" msg += "```" await ctx.channel.send(msg) # @profile_group.command('force_set_level') # @commands.check(isWhiteList) # async def profile_force_set_level(self, ctx:commands.Context, rank=10): # message = ctx.message # member: Member = self.db.get_member(message.author.id) # if not member: # member = self.db.add_member(message.author.id) # if member.rank != rank: # self.db.force_update_rank(member.member_id, rank) # channel = ChannelUtil.getMessageChannelId(message.guild.id) # if channel == None: # channelToUse = message.channel # else: # channelToUse = message.guild.get_channel(channel) # await channelToUse.send('恭喜<@{}> 等級提升至{}。'.format(message.author.id, rank)) # await self.updateUserKfpRoles(message, rank, channelToUse) # @profile_group.command('force_set_exp') # async def profile_force_set_exp(self, ctx:commands.Context, exp=-1): # if exp < 1: # return # message = ctx.message # member: Member = self.db.get_member(message.author.id) # if not member: # member = self.db.add_member(message.author.id) # self.db.set_exp(member.member_id, exp) # member: Member = self.db.get_member(message.author.id) # await ctx.send(f"你現在有經驗值:{member.exp}") # @profile_group.command('get_rank_exp') # async def profile_get_rank_exp(self, ctx:commands.Context, rank=0): # if rank < 1: # return # exp_needed = Util.get_rank_exp(rank) # await ctx.send(f"要升級到等級{rank} 需要經驗值:{exp_needed}") @profile_group.command('items') async def show_items_command(self, ctx: commands.Context): if ctx.author.bot: return # ignore bot records = InventoryUtil.getAllItemsBelongToUser( ctx.guild.id, ctx.author.id) msg = "" if len(records) > 0: msg += "你現在有以下物品:\n" record: InventoryRecord for record in records: msg += f"{record.item.name} x {record.amount}\n" else: msg += "你目前沒有任何物品" await ctx.author.send(msg) @profile_group.command('syncAllRank') @commands.check(isWhiteList) async def reset_everyone_rank(self, ctx: commands.Context, rank=0): member_id_list = Member.select( Member.member_id).where(Member.rank >= rank) for member_id in member_id_list: member: Member = self.db.get_member(member_id) try: user = await ctx.guild.fetch_member(member_id) except NotFound as e: continue await self.__updateUserRole(ctx.guild, user, member, member.rank, None, True) @profile_group.command(name='leaderboard') @commands.check(isWhiteList) async def profile_leaderboard(self, ctx: commands.Context, limit=10): max_limit = 25 if limit > max_limit: await ctx.channel.send(f'{limit} 超過上限, 請選擇小於 {max_limit} 的數字') return top_leaders = self.db.get_leader_board(limit) msg = "```" msg += "員工等級排名:\n" member: Member for rank, member in enumerate(top_leaders): guild_member = ctx.guild.get_member(member.member_id) user = await self.bot.fetch_user(member.member_id) if guild_member: if guild_member.nick: msg += f"第{rank+1}名: {guild_member.nick}\n" else: msg += f"第{rank+1}名: {guild_member.display_name}\n" elif user: msg += f"第{rank+1}名: {user.display_name}\n" msg += "```" await ctx.channel.send(msg) async def updateUserKfpRoles(self, message: Message, rank: int, channelToUse: GuildChannel): member = Member.select().where(Member.member_id == message.author.id) user = message.author await self.__updateUserRole(message.guild, user, member, rank, channelToUse, False) async def __updateUserRole(self, guild: Guild, user: User, member: Member, rank: int, channelToUse: GuildChannel, internal: bool): if user: if member: newRoles = RoleUtil.getKfpRolesFromLevel(guild.id, rank) if len(newRoles) > 0: for newRole in newRoles: newGuildRole: Role = guild.get_role(newRole.role_id) if newGuildRole: if not newGuildRole in user.roles: # 用戶有新身份組 # 先移除所有不符合的身份組 oldRoles: KfpRole = RoleUtil.getCurrentRoles( guild.id, Util.RoleCategory(newRole.category)) if oldRoles: oldGuildRoles = [] for oldRole in oldRoles: guildRole = guild.get_role( oldRole.role_id) if guildRole and guildRole in user.roles: oldGuildRoles.append(guildRole) for oldGuildRole in oldGuildRoles: await user.remove_roles(oldGuildRole) # 添加新的身份組 await user.add_roles(newGuildRole) if internal: print( "adding role {} to member {} successed!" .format(newGuildRole.name, user.name)) else: embed = Embed() embed.description = '恭喜<@!{}> 成為 {}'.format( user.id, newGuildRole.name) await channelToUse.send(embed=embed) def populateChannels(self, message: Message, isTest: bool): if isTest: return False if len(self.__channels) == 0: categories = message.guild.categories for category in categories: if category.name == "🐔員工大廳-Hühnerfarm": channels = category.channels result = [] for channel in channels: result.append(channel.id) self.__channels = result return True return False def channelAllowed(self, channel_id: int, isTest: bool): if isTest: return True return channel_id in self.__channels