From ff9558586ec866de2ad8dda51e0292893db6a49e Mon Sep 17 00:00:00 2001 From: pushiscool Date: Mon, 25 May 2026 03:22:45 -0400 Subject: [PATCH 1/2] Fix study role expiration and DM --- src/cogs/moderation/infraction.py | 8 ++- src/cogs/study.py | 116 ++++++++++++++++++++++++------ tests/test_study_helpers.py | 42 +++++++++++ 3 files changed, 144 insertions(+), 22 deletions(-) create mode 100644 tests/test_study_helpers.py diff --git a/src/cogs/moderation/infraction.py b/src/cogs/moderation/infraction.py index d6c4f1c..c7d1096 100644 --- a/src/cogs/moderation/infraction.py +++ b/src/cogs/moderation/infraction.py @@ -30,8 +30,12 @@ def __init__(self, bot: APBot) -> None: def has_mod_role(self, member: Member) -> bool: allowed_role_names = {"Trial Chat Moderator", "Chat Moderator", "Admin"} allowed_role_ids = { - conf.get("bot_staff_role_id"), - conf.get("special_perms_role_id"), + role_id + for role_id in ( + conf.get("bot_staff_role_id"), + conf.get("special_perms_role_id"), + ) + if role_id is not None } perms = getattr(member, "guild_permissions", None) diff --git a/src/cogs/study.py b/src/cogs/study.py index 5d998bb..00cfb57 100644 --- a/src/cogs/study.py +++ b/src/cogs/study.py @@ -4,15 +4,52 @@ import nextcord from nextcord import Interaction, Embed, Attachment, Message, slash_command, SlashOption -from nextcord.ext import commands +from nextcord.ext import commands, tasks, tasks from typing import Union from cogs.utils import convert_time blue = nextcord.Color.teal() -from config_handler import Config -config_path = "config.json" -conf = Config(config_path) +from app_config import get_command_guild_ids, load_optional_config +conf = load_optional_config() +COMMAND_GUILD_IDS = get_command_guild_ids(conf) + +MIN_STUDY_SECONDS = 60 * 10 +MAX_STUDY_SECONDS = 60 * 60 * 24 * 7 + + +def validate_study_duration(duration: str): + duration_seconds = convert_time(duration) + + if isinstance(duration_seconds, str): + return None, duration_seconds + + if duration_seconds < MIN_STUDY_SECONDS: + return None, "Please set a duration greater than or equal to 10 minutes." + + if duration_seconds > MAX_STUDY_SECONDS: + return None, "Please set a duration less than or equal to one week." + + return duration_seconds, None + + +def get_study_delay_seconds(study_end: int, now: int | None = None) -> int: + current_time = int(time.time()) if now is None else now + return max(0, int(study_end) - current_time) + + +def get_study_delay_seconds(study_end: int, now: int | None = None) -> int: + current_time = int(time.time()) if now is None else now + return max(0, int(study_end) - current_time) + + +async def send_study_role_dm(member, study_end: int) -> None: + try: + await member.send( + f"You enabled Study Mode. The Study role will be removed at ." + ) + except (nextcord.Forbidden, nextcord.HTTPException): + return class QuestionConfirm(nextcord.ui.View): @@ -53,7 +90,7 @@ def __init__(self, bot: commands.Bot): self.potd_cooldowns = {} - @slash_command(name="question", description="Ask a question in a subject channel.", guild_ids=[conf.get("guild_id")]) + @slash_command(name="question", description="Ask a question in a subject channel.", guild_ids=COMMAND_GUILD_IDS) async def question( self, inter: Interaction, @@ -106,7 +143,7 @@ async def question( - @slash_command(name="potd", description="Make a problem-of-the-day for a subject channel.", guild_ids=[conf.get("guild_id")]) + @slash_command(name="potd", description="Make a problem-of-the-day for a subject channel.", guild_ids=COMMAND_GUILD_IDS) async def potd( self, inter: Interaction, @@ -174,6 +211,46 @@ async def potd( + @tasks.loop(minutes=1) + async def expire_study_roles(self) -> None: + try: + study_users = await self.bot.db.study.get_all() + except Exception: + return + + now = int(time.time()) + for user_id, study_end in study_users.items(): + try: + if int(study_end) <= now: + await self.remove_study_role(int(user_id)) + except Exception: + continue + + @commands.Cog.listener() + async def on_ready(self) -> None: + if not self.expire_study_roles.is_running(): + self.expire_study_roles.start() + + @tasks.loop(minutes=1) + async def expire_study_roles(self) -> None: + try: + study_users = await self.bot.db.study.get_all() + except Exception: + return + + now = int(time.time()) + for user_id, study_end in study_users.items(): + try: + if int(study_end) <= now: + await self.remove_study_role(int(user_id)) + except Exception: + continue + + @commands.Cog.listener() + async def on_ready(self) -> None: + if not self.expire_study_roles.is_running(): + self.expire_study_roles.start() + async def remove_study_role(self, user_id: int) -> None: try: guild = await self.bot.fetch_guild(self.bot.config.get("guild_id")) @@ -188,7 +265,7 @@ async def remove_study_role(self, user_id: int) -> None: await self.bot.db.study.delete_user(user_id) await user.remove_roles(study_role) - @slash_command(name="study", description="Prevent yourself from viewing unhelpful channels.", guild_ids=[conf.get("guild_id")]) + @slash_command(name="study", description="Prevent yourself from viewing unhelpful channels.", guild_ids=COMMAND_GUILD_IDS) async def study( self, inter: Interaction, @@ -197,32 +274,31 @@ async def study( await inter.response.defer(ephemeral=True) resp = await inter.send("Validating time input...", ephemeral=True) - duration: Union[str, int] = convert_time(duration) - if isinstance(duration, str): - return await resp.edit(content=duration) - - if duration < 60 * 10: - return await resp.edit("Please set a duration greater than 10 minutes!") - if duration > 60 * 60 * 24 * 7: - return await resp.edit("Please set a duration lesser than a week") + duration_seconds, error = validate_study_duration(duration) + if error: + return await resp.edit(content=error) - await resp.edit("Performing actions...") - study_end = int(time.time()) + duration + await resp.edit(content="Performing actions...") + study_end = int(time.time()) + duration_seconds guild = await self.bot.fetch_guild(self.bot.config.get("guild_id")) study_role = nextcord.utils.get(guild.roles, name="Study") + if study_role is None: + return await resp.edit(content="Study role not found.") + if study_role.id in [role.id for role in inter.user.roles]: await resp.edit("Removing role...") await self.remove_study_role(inter.user.id) await inter.user.add_roles(study_role) - await resp.edit("Updating database...") + await resp.edit(content="Updating database...") await self.bot.db.study.set_time(inter.user.id, study_end) - self.bot.loop.call_later(duration, asyncio.create_task, self.remove_study_role(inter.user.id)) + self.bot.loop.call_later(get_study_delay_seconds(study_end), asyncio.create_task, self.remove_study_role(inter.user.id)) + await send_study_role_dm(inter.user, study_end) await resp.edit( - f"The study role will be removed at ." + content=f"The study role will be removed at ." ) class ConfirmDeny(nextcord.ui.View): diff --git a/tests/test_study_helpers.py b/tests/test_study_helpers.py new file mode 100644 index 0000000..52aaff4 --- /dev/null +++ b/tests/test_study_helpers.py @@ -0,0 +1,42 @@ +import asyncio + +from cogs.study import MAX_STUDY_SECONDS, MIN_STUDY_SECONDS, get_study_delay_seconds, send_study_role_dm, validate_study_duration + + +class FakeMember: + def __init__(self): + self.sent = None + + async def send(self, content): + self.sent = content + + +def test_validate_study_duration_accepts_bounds(): + assert validate_study_duration("10m") == (MIN_STUDY_SECONDS, None) + assert validate_study_duration("168h") == (MAX_STUDY_SECONDS, None) + + +def test_validate_study_duration_rejects_out_of_range(): + assert validate_study_duration("9m")[0] is None + assert validate_study_duration("169h")[0] is None + + +def test_send_study_role_dm_sends_message(): + member = FakeMember() + + asyncio.run(send_study_role_dm(member, 1234567890)) + + assert "Study Mode" in member.sent + assert "" in member.sent + + +def test_get_study_delay_seconds(): + assert get_study_delay_seconds(1000, now=900) == 100 + assert get_study_delay_seconds(1000, now=1000) == 0 + assert get_study_delay_seconds(1000, now=1100) == 0 + + +def test_get_study_delay_seconds(): + assert get_study_delay_seconds(1000, now=900) == 100 + assert get_study_delay_seconds(1000, now=1000) == 0 + assert get_study_delay_seconds(1000, now=1100) == 0 From 23741e07927e5b37a4501a3b174b1c40267cc78e Mon Sep 17 00:00:00 2001 From: pushiscool Date: Mon, 25 May 2026 15:19:13 -0400 Subject: [PATCH 2/2] Harden study role expiration --- src/cogs/study.py | 83 +++++++++++++++++++------------------ tests/test_study_helpers.py | 10 ++--- 2 files changed, 45 insertions(+), 48 deletions(-) diff --git a/src/cogs/study.py b/src/cogs/study.py index 00cfb57..7eff8ce 100644 --- a/src/cogs/study.py +++ b/src/cogs/study.py @@ -3,9 +3,8 @@ import asyncio import nextcord -from nextcord import Interaction, Embed, Attachment, Message, slash_command, SlashOption -from nextcord.ext import commands, tasks, tasks -from typing import Union +from nextcord import Interaction, Embed, Attachment, slash_command, SlashOption +from nextcord.ext import commands, tasks from cogs.utils import convert_time @@ -38,18 +37,15 @@ def get_study_delay_seconds(study_end: int, now: int | None = None) -> int: return max(0, int(study_end) - current_time) -def get_study_delay_seconds(study_end: int, now: int | None = None) -> int: - current_time = int(time.time()) if now is None else now - return max(0, int(study_end) - current_time) - - -async def send_study_role_dm(member, study_end: int) -> None: +async def send_study_role_dm(member, study_end: int) -> bool: try: await member.send( f"You enabled Study Mode. The Study role will be removed at ." ) except (nextcord.Forbidden, nextcord.HTTPException): - return + return False + + return True class QuestionConfirm(nextcord.ui.View): @@ -88,6 +84,10 @@ def __init__(self, bot: commands.Bot): self.bot = bot self.last_topic_update = {} self.potd_cooldowns = {} + + def cog_unload(self) -> None: + if self.expire_study_roles.is_running(): + self.expire_study_roles.cancel() @slash_command(name="question", description="Ask a question in a subject channel.", guild_ids=COMMAND_GUILD_IDS) @@ -215,7 +215,8 @@ async def potd( async def expire_study_roles(self) -> None: try: study_users = await self.bot.db.study.get_all() - except Exception: + except Exception as exc: + print(f"[Study] Failed to fetch study role expirations: {exc}") return now = int(time.time()) @@ -226,25 +227,9 @@ async def expire_study_roles(self) -> None: except Exception: continue - @commands.Cog.listener() - async def on_ready(self) -> None: - if not self.expire_study_roles.is_running(): - self.expire_study_roles.start() - - @tasks.loop(minutes=1) - async def expire_study_roles(self) -> None: - try: - study_users = await self.bot.db.study.get_all() - except Exception: - return - - now = int(time.time()) - for user_id, study_end in study_users.items(): - try: - if int(study_end) <= now: - await self.remove_study_role(int(user_id)) - except Exception: - continue + @expire_study_roles.before_loop + async def before_expire_study_roles(self) -> None: + await self.bot.wait_until_ready() @commands.Cog.listener() async def on_ready(self) -> None: @@ -252,18 +237,29 @@ async def on_ready(self) -> None: self.expire_study_roles.start() async def remove_study_role(self, user_id: int) -> None: + guild_id = self.bot.config.get("guild_id") + try: - guild = await self.bot.fetch_guild(self.bot.config.get("guild_id")) + guild = await self.bot.getch_guild(guild_id) user = await self.bot.getch_member(guild.id, user_id) study_role = nextcord.utils.get(guild.roles, name="Study") - except: + except Exception: return - if not user or not study_role: + if not study_role: return + if not user: + await self.bot.db.study.delete_user(user_id) + return + + if study_role.id in [role.id for role in user.roles]: + try: + await user.remove_roles(study_role, reason="Study role duration expired.") + except (nextcord.Forbidden, nextcord.HTTPException): + return + await self.bot.db.study.delete_user(user_id) - await user.remove_roles(study_role) @slash_command(name="study", description="Prevent yourself from viewing unhelpful channels.", guild_ids=COMMAND_GUILD_IDS) async def study( @@ -281,21 +277,26 @@ async def study( await resp.edit(content="Performing actions...") study_end = int(time.time()) + duration_seconds - guild = await self.bot.fetch_guild(self.bot.config.get("guild_id")) + guild = await self.bot.getch_guild(self.bot.config.get("guild_id")) + if not guild: + return await resp.edit(content="Server not found.") + study_role = nextcord.utils.get(guild.roles, name="Study") if study_role is None: return await resp.edit(content="Study role not found.") - if study_role.id in [role.id for role in inter.user.roles]: - await resp.edit("Removing role...") - await self.remove_study_role(inter.user.id) - - await inter.user.add_roles(study_role) + user_role_ids = {role.id for role in inter.user.roles} + if study_role.id not in user_role_ids: + await inter.user.add_roles(study_role, reason="Study role enabled by user.") await resp.edit(content="Updating database...") await self.bot.db.study.set_time(inter.user.id, study_end) - self.bot.loop.call_later(get_study_delay_seconds(study_end), asyncio.create_task, self.remove_study_role(inter.user.id)) + user_id = inter.user.id + self.bot.loop.call_later( + get_study_delay_seconds(study_end), + lambda user_id=user_id: asyncio.create_task(self.remove_study_role(user_id)), + ) await send_study_role_dm(inter.user, study_end) await resp.edit( content=f"The study role will be removed at ." diff --git a/tests/test_study_helpers.py b/tests/test_study_helpers.py index 52aaff4..b13767f 100644 --- a/tests/test_study_helpers.py +++ b/tests/test_study_helpers.py @@ -14,6 +14,7 @@ async def send(self, content): def test_validate_study_duration_accepts_bounds(): assert validate_study_duration("10m") == (MIN_STUDY_SECONDS, None) assert validate_study_duration("168h") == (MAX_STUDY_SECONDS, None) + assert validate_study_duration("1w") == (MAX_STUDY_SECONDS, None) def test_validate_study_duration_rejects_out_of_range(): @@ -24,8 +25,9 @@ def test_validate_study_duration_rejects_out_of_range(): def test_send_study_role_dm_sends_message(): member = FakeMember() - asyncio.run(send_study_role_dm(member, 1234567890)) + result = asyncio.run(send_study_role_dm(member, 1234567890)) + assert result is True assert "Study Mode" in member.sent assert "" in member.sent @@ -34,9 +36,3 @@ def test_get_study_delay_seconds(): assert get_study_delay_seconds(1000, now=900) == 100 assert get_study_delay_seconds(1000, now=1000) == 0 assert get_study_delay_seconds(1000, now=1100) == 0 - - -def test_get_study_delay_seconds(): - assert get_study_delay_seconds(1000, now=900) == 100 - assert get_study_delay_seconds(1000, now=1000) == 0 - assert get_study_delay_seconds(1000, now=1100) == 0