basic implementation of send_after

Author: artcz

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,18 +69,21 @@ 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)
     else:
         # 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
 
 

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):

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()