diff --git a/intbot/core/bot/main.py b/intbot/core/bot/main.py index c6e8584..feeba54 100644 --- a/intbot/core/bot/main.py +++ b/intbot/core/bot/main.py @@ -2,6 +2,7 @@ from core.models import DiscordMessage from discord.ext import commands, tasks from django.conf import settings +from django.db.models import Q from django.utils import timezone intents = discord.Intents.default() @@ -38,19 +39,22 @@ async def wiki(ctx): suppress_embeds=True, ) + @bot.command() async def close(ctx): channel = ctx.channel author = ctx.message.author # Check if it's a public or private post (thread) - if channel.type in (discord.ChannelType.public_thread, discord.ChannelType.private_thread): + if channel.type in ( + discord.ChannelType.public_thread, + discord.ChannelType.private_thread, + ): parent = channel.parent # Check if the post (thread) was sent in a forum, # so we can add a tag if parent.type == discord.ChannelType.forum: - # Get tag from forum tag = None for _tag in parent.available_tags: @@ -65,7 +69,9 @@ async def close(ctx): await ctx.message.delete() # Send notification to the thread - await channel.send(f"# This was marked as done by {author.mention}", suppress_embeds=True) + await channel.send( + f"# This was marked as done by {author.mention}", suppress_embeds=True + ) # We need to archive after adding tags in case it was a forum. await channel.edit(archived=True) @@ -73,10 +79,11 @@ async def close(ctx): # Remove command message await ctx.message.delete() - await channel.send("The !close command is intended to be used inside a thread/post", - suppress_embeds=True, - delete_after=5) - + await channel.send( + "The !close command is intended to be used inside a thread/post", + suppress_embeds=True, + delete_after=5, + ) @bot.command() @@ -94,7 +101,10 @@ async def qlen(ctx): def get_messages(): - messages = DiscordMessage.objects.filter(sent_at__isnull=True) + messages = DiscordMessage.objects.filter( + Q(send_after__isnull=True) | Q(send_after__lte=timezone.now()), + sent_at__isnull=True, + ) return messages diff --git a/intbot/core/migrations/0004_add_send_after.py b/intbot/core/migrations/0004_add_send_after.py new file mode 100644 index 0000000..d361823 --- /dev/null +++ b/intbot/core/migrations/0004_add_send_after.py @@ -0,0 +1,18 @@ +# Generated by Django 5.1.4 on 2025-02-27 00:50 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('core', '0003_added_extra_field_to_webhook'), + ] + + operations = [ + migrations.AddField( + model_name='discordmessage', + name='send_after', + field=models.DateTimeField(blank=True, null=True), + ), + ] diff --git a/intbot/core/models.py b/intbot/core/models.py index bdf9df3..d6a36c8 100644 --- a/intbot/core/models.py +++ b/intbot/core/models.py @@ -44,7 +44,10 @@ class DiscordMessage(models.Model): created_at = models.DateTimeField(auto_now_add=True) modified_at = models.DateTimeField(auto_now=True) - # Messages to be have null here + # To delay messages to be sent after certain timestmap in the future + send_after = models.DateTimeField(blank=True, null=True) + + # Messages to be sent have null here sent_at = models.DateTimeField(blank=True, null=True) def __str__(self): diff --git a/intbot/tests/test_bot/test_main.py b/intbot/tests/test_bot/test_main.py index ea3517e..1afeadf 100644 --- a/intbot/tests/test_bot/test_main.py +++ b/intbot/tests/test_bot/test_main.py @@ -1,14 +1,14 @@ +import contextlib +from datetime import timedelta from unittest import mock from unittest.mock import AsyncMock, patch -import contextlib -import discord - -from django.db import connections +import discord import pytest from asgiref.sync import sync_to_async -from core.bot.main import ping, poll_database, qlen, source, version, wiki, close +from core.bot.main import close, ping, poll_database, qlen, source, version, wiki from core.models import DiscordMessage +from django.db import connections from django.utils import timezone # NOTE(artcz) @@ -101,6 +101,7 @@ async def test_wiki_command(): suppress_embeds=True, ) + @pytest.mark.asyncio async def test_close_command_working(): # Mock context @@ -118,6 +119,7 @@ async def test_close_command_working(): suppress_embeds=True, ) + @pytest.mark.asyncio async def test_close_command_notworking(): # Mock context @@ -131,7 +133,7 @@ async def test_close_command_notworking(): ctx.channel.send.assert_called_once_with( "The !close command is intended to be used inside a thread/post", suppress_embeds=True, - delete_after=5 + delete_after=5, ) @@ -235,12 +237,40 @@ async def test_polling_messages_sends_nothing_if_all_messages_are_sent(): @pytest.mark.asyncio @pytest.mark.django_db -async def test_polling_messages_sends_message_if_not_sent_and_sets_sent_at(): +async def test_polling_messages_sends_nothing_if_all_messages_in_the_future(): + mock_channel = AsyncMock() + mock_channel.send = AsyncMock() + await DiscordMessage.objects.acreate( + send_after=timezone.now() + timedelta(hours=3), + sent_at=None, + ) + + with patch("core.bot.main.bot.get_channel", return_value=mock_channel): + await poll_database() + + mock_channel.send.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.django_db +@pytest.mark.parametrize( + "send_after", + [ + None, + timezone.now(), + ], + ids=[ + "send_after_isnull", + "send_after_is_in_the_past", + ], +) +async def test_polling_messages_sends_message_if_not_sent_and_sets_sent_at(send_after): start = timezone.now() dm = await DiscordMessage.objects.acreate( channel_id="1234", content="asdf", sent_at=None, + send_after=send_after, ) mock_channel = AsyncMock() mock_channel.send = AsyncMock()