diff --git a/src/cmds/automation/scheduled_tasks.py b/src/cmds/automation/scheduled_tasks.py index ceed6fd..9b763fc 100644 --- a/src/cmds/automation/scheduled_tasks.py +++ b/src/cmds/automation/scheduled_tasks.py @@ -43,22 +43,25 @@ async def auto_unban(self) -> None: bans = result.all() logger.debug(f"Got {len(bans)} bans from DB.") - for ban in bans: - run_at = datetime.fromtimestamp(ban.unban_time) - logger.debug( - f"Got user_id: {ban.user_id} and unban timestamp: {run_at} from DB." - ) - - for guild_id in settings.guild_ids: - logger.debug(f"Running for guild: {guild_id}.") - guild = self.bot.get_guild(guild_id) - if guild: - member = await self.bot.get_member_or_user(guild, ban.user_id) - unban_task = schedule(unban_member(guild, member), run_at=run_at) - unban_tasks.append(unban_task) - logger.info(f"Scheduled unban task for user_id {ban.user_id} at {run_at}.") - else: - logger.warning(f"Unable to find guild with ID {guild_id}.") + for guild_id in settings.guild_ids: + logger.debug(f"Running for guild: {guild_id}.") + guild = self.bot.get_guild(guild_id) + if not guild: + logger.warning(f"Unable to find guild with ID {guild_id}.") + continue + + for ban in bans: + run_at = datetime.fromtimestamp(ban.unban_time) + logger.debug( + f"Got user_id: {ban.user_id} and unban timestamp: {run_at} from DB." + ) + member = await self.bot.get_member_or_user(guild, ban.user_id) + if not member: + logger.info(f"Member with id: {ban.user_id} not found.") + continue + unban_task = schedule(unban_member(guild, member), run_at=run_at) + unban_tasks.append(unban_task) + logger.info(f"Scheduled unban task for user_id {ban.user_id} at {run_at}.") await asyncio.gather(*unban_tasks) @@ -72,23 +75,27 @@ async def auto_unmute(self) -> None: mutes = result.all() logger.debug(f"Got {len(mutes)} mutes from DB.") - for mute in mutes: - run_at = datetime.fromtimestamp(mute.unmute_time) - logger.debug( - "Got user_id: {user_id} and unmute timestamp: {unmute_ts} from DB.".format( - user_id=mute.user_id, unmute_ts=run_at + for guild_id in settings.guild_ids: + guild = self.bot.get_guild(guild_id) + if not guild: + logger.warning(f"Unable to find guild with ID {guild_id}.") + continue + + for mute in mutes: + run_at = datetime.fromtimestamp(mute.unmute_time) + logger.debug( + "Got user_id: {user_id} and unmute timestamp: {unmute_ts} from DB.".format( + user_id=mute.user_id, unmute_ts=run_at + ) ) - ) + member = await self.bot.get_member_or_user(guild, mute.user_id) + if not member: + logger.info(f"Member with id: {mute.user_id} not found.") + continue + unmute_task = schedule(unmute_member(guild, member), run_at=run_at) + unmute_tasks.append(unmute_task) + logger.info(f"Scheduled unban task for user_id {mute.user_id} at {str(run_at)}.") - for guild_id in settings.guild_ids: - guild = self.bot.get_guild(guild_id) - if guild: - member = await self.bot.get_member_or_user(guild, mute.user_id) - unmute_task = schedule(unmute_member(guild, member), run_at=run_at) - unmute_tasks.append(unmute_task) - logger.info(f"Scheduled unban task for user_id {mute.user_id} at {str(run_at)}.") - else: - logger.warning(f"Unable to find guild with ID {guild_id}.") await asyncio.gather(*unmute_tasks) diff --git a/tests/src/cmds/core/test_ban.py b/tests/src/cmds/core/test_ban.py index e8073f8..8308ff7 100644 --- a/tests/src/cmds/core/test_ban.py +++ b/tests/src/cmds/core/test_ban.py @@ -21,8 +21,10 @@ async def test_ban_success(self, ctx, bot): user = helpers.MockMember(id=2, name="Banned User") bot.get_member_or_user.return_value = user - with patch('src.cmds.core.ban.ban_member', new_callable=AsyncMock) as ban_member_mock, \ - patch('src.cmds.core.ban.add_evidence_note', new_callable=AsyncMock) as add_evidence_note_mock: + with ( + patch('src.cmds.core.ban.ban_member', new_callable=AsyncMock) as ban_member_mock, + patch('src.cmds.core.ban.add_evidence_note', new_callable=AsyncMock) as add_evidence_note_mock + ): ban_response = SimpleResponse( message=f"Member {user.display_name} has been banned permanently.", delete_after=0 ) @@ -68,15 +70,16 @@ async def test_tempban_success(self, ctx, bot): ) @pytest.mark.asyncio - @pytest.mark.skip(reason="Skipping test temporarily until figuring out what's going wrong.") - async def test_tempban_failed_with_wrong_duration(self, ctx, bot): + async def test_tempban_failed_with_wrong_duration(self, ctx, bot, guild): ctx.user = helpers.MockMember(id=1, name="Test User") + ctx.guild = guild user = helpers.MockMember(id=2, name="Banned User") bot.get_member_or_user.return_value = user with ( patch('src.helpers.ban.validate_duration', new_callable=AsyncMock) as validate_duration_mock, - patch('src.cmds.core.ban.ban_member', new_callable=AsyncMock) as ban_member_mock + patch('src.cmds.core.ban.ban_member', new_callable=AsyncMock) as ban_member_mock, + patch('src.cmds.core.ban.add_evidence_note', new_callable=AsyncMock) as add_evidence_note_mock, ): validate_duration_mock.return_value = ( 0, "Malformed duration. Please use duration units, (e.g. 12h, 14d, 5w)." @@ -90,6 +93,9 @@ async def test_tempban_failed_with_wrong_duration(self, ctx, bot): await cog.tempban.callback(cog, ctx, user, "5", "Any valid reason", "Some evidence") # Assertions + add_evidence_note_mock.assert_called_once_with( + user.id, "ban", "Any valid reason", "Some evidence", ctx.user.id + ) ban_member_mock.assert_called_once_with( bot, ctx.guild, user, "5", "Any valid reason", "Some evidence", ctx.user, needs_approval=True )