added support for scheduling messages, including debugging async tests

Author: artcz

Makefile

@@ -49,7 +49,14 @@ bot:
 # ================
 
 test:
-	$(TEST_CMD) -s -vv
+	$(TEST_CMD) -s -v
+
+test/k:
+	$(TEST_CMD) -s -v -k $(K)
+
+test/fast:
+	# skip slow tests
+	$(TEST_CMD) -s -v -m "not slow"
 
 
 lint:

intbot/core/bot/main.py

@@ -1,6 +1,8 @@
 import discord
+from core.models import DiscordMessage
+from discord.ext import commands, tasks
 from django.conf import settings
-from discord.ext import commands
+from django.utils import timezone
 
 intents = discord.Intents.default()
 intents.members = True
@@ -33,6 +35,38 @@ async def version(ctx):
     await ctx.send(f"Version: {app_version}")
 
 
+@bot.command()
+async def qlen(ctx):
+    qs = get_messages()
+    # qs = DiscordMessage.objects.afilter(sent_at__isnull=True)
+    cnt = await qs.acount()
+    await ctx.send(f"In the queue there are: {cnt} messages")
+
+
+def get_messages():
+    messages = DiscordMessage.objects.filter(sent_at__isnull=True)
+    return messages
+
+
+@tasks.loop(seconds=60)  # Seconds
+async def poll_database():
+    """Check for unsent messages and send them."""
+    messages = get_messages()
+    print("Polling database.... ", timezone.now())
+
+    async for message in messages:
+        channel = bot.get_channel(int(message.channel_id))
+        if channel:
+            await channel.send(
+                message.content,
+                suppress_embeds=True,
+            )
+            message.sent_at = timezone.now()
+            await message.asave()
+        else:
+            print("Channel does not exist!")
+
+
 def run_bot():
     bot_token = settings.DISCORD_BOT_TOKEN
     bot.run(bot_token)

intbot/intbot/settings.py

@@ -148,7 +148,7 @@
     DATABASES = {
         "default": {
             "ENGINE": "django.db.backends.postgresql",
-            "NAME": "intbot_database_dev",
+            "NAME": "intbot_database_test",
             "USER": "intbot_user",
             "PASSWORD": "intbot_password",
             "HOST": "localhost",

intbot/tests/test_bot.py

@@ -1,7 +1,72 @@
-from unittest.mock import AsyncMock
+from unittest import mock
+from unittest.mock import AsyncMock, patch
+import contextlib
+
+from django.db import connections
 
 import pytest
-from core.bot.main import ping, version, source
+from asgiref.sync import sync_to_async
+from core.bot.main import ping, poll_database, qlen, source, version
+from core.models import DiscordMessage
+from django.utils import timezone
+
+# NOTE(artcz)
+# The fixture below (fix_async_db) is copied from this issue
+# https://github.com/pytest-dev/pytest-asyncio/issues/226
+# it seems to fix the issue and also speed up the test from ~6s down to 1s.
+# Thanks to (@gbdlin) for help with debugging.
+
+@pytest.fixture(autouse=True)
+def fix_async_db(request):
+    """
+
+    If you don't use this fixture for async tests that use the ORM/database
+    you won't get proper teardown of the database.
+    This is a bug somehwere in pytest-django, pytest-asyncio or django itself.
+
+    Nobody knows how to solve it, or who should solve it.
+    Workaround here: https://github.com/django/channels/issues/1091#issuecomment-701361358
+    More info:
+    https://github.com/pytest-dev/pytest-django/issues/580
+    https://code.djangoproject.com/ticket/32409
+    https://github.com/pytest-dev/pytest-asyncio/issues/226
+
+
+    The actual implementation of this workaround constists on ensuring
+    Django always returns the same database connection independently of the thread
+    the code that requests a db connection is in.
+
+    We were unable to use better patching methods (the target is asgiref/local.py),
+    so we resorted to mocking the _lock_storage context manager so that it returns a Mock.
+    That mock contains the default connection of the main thread (instead of the connection
+    of the running thread).
+
+    This only works because our tests only ever use the default connection, which is the only thing our mock returns.
+
+    We apologize in advance for the shitty implementation.
+    """
+    if request.node.get_closest_marker("asyncio") is None or request.node.get_closest_marker("django_db") is None:
+        # Only run for async tests that use the database
+        yield
+        return
+
+    main_thread_local = connections._connections
+    for conn in connections.all():
+        conn.inc_thread_sharing()
+
+    main_thread_default_conn = main_thread_local._storage.default
+    main_thread_storage = main_thread_local._lock_storage
+
+    @contextlib.contextmanager
+    def _lock_storage():
+        yield mock.Mock(default=main_thread_default_conn)
+
+    try:
+        with patch.object(main_thread_default_conn, "close"):
+            object.__setattr__(main_thread_local, "_lock_storage", _lock_storage)
+            yield
+    finally:
+        object.__setattr__(main_thread_local, "_lock_storage", main_thread_storage)
 
 
 @pytest.mark.asyncio
@@ -41,3 +106,105 @@ async def test_source_command():
         "I'm here: https://github.com/europython/internal-bot",
         suppress_embeds=True,
     )
+
+
+@pytest.mark.asyncio
+@pytest.mark.django_db
+async def test_qlen_command_returns_zero_if_no_messages():
+    # Mock context
+    ctx = AsyncMock()
+
+    # Call the command
+    await qlen(ctx)
+
+    # Assert that the command sent the expected message
+    ctx.send.assert_called_once_with("In the queue there are: 0 messages")
+
+
+@pytest.mark.asyncio
+@pytest.mark.slow
+@pytest.mark.django_db
+async def test_qlen_command_returns_zero_if_all_messages_sent():
+    # Mock context
+    ctx = AsyncMock()
+    await DiscordMessage.objects.acreate(sent_at=timezone.now())
+
+    # Call the command
+    await qlen(ctx)
+
+    # Assert that the command sent the expected message
+    ctx.send.assert_called_once_with("In the queue there are: 0 messages")
+
+
+@pytest.mark.asyncio
+@pytest.mark.slow
+@pytest.mark.django_db
+async def test_qlen_command_correctly_counts_unsent_messags():
+    # Mock context
+    ctx = AsyncMock()
+    for _ in range(3):
+        await DiscordMessage.objects.acreate(
+            channel_id="1234",
+            content="foobar",
+            sent_at=None,
+        )
+
+    # Call the command
+    await qlen(ctx)
+
+    # Assert that the command sent the expected message
+    ctx.send.assert_called_once_with("In the queue there are: 3 messages")
+
+
+@pytest.mark.asyncio
+@pytest.mark.slow
+@pytest.mark.django_db
+async def test_polling_messages_sends_nothing_without_messages():
+    # NOTE: For some reason this test slows down the testsuite a bit
+    # breakpoint()
+    mock_channel = AsyncMock()
+    mock_channel.send = AsyncMock()
+
+    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.slow
+@pytest.mark.django_db
+async def test_polling_messages_sends_nothing_if_all_messages_are_sent():
+    # NOTE: For some reason this test slows down the testsuite a bit
+    mock_channel = AsyncMock()
+    mock_channel.send = AsyncMock()
+    await DiscordMessage.objects.acreate(sent_at=timezone.now())
+
+    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.slow
+@pytest.mark.django_db
+async def test_polling_messages_sends_message_if_not_sent_and_sets_sent_at():
+    # NOTE: For some reason this test slows down the testsuite a bit
+    start = timezone.now()
+    dm = await DiscordMessage.objects.acreate(
+        channel_id="1234",
+        content="asdf",
+        sent_at=None,
+    )
+    mock_channel = AsyncMock()
+    mock_channel.send = AsyncMock()
+
+    with patch("core.bot.main.bot.get_channel", return_value=mock_channel):
+        await poll_database()
+
+    mock_channel.send.assert_called_once_with("asdf", suppress_embeds=True)
+    await sync_to_async(dm.refresh_from_db)()
+    assert dm.sent_at is not None
+    end = timezone.now()
+    assert start < dm.sent_at < end

pyproject.toml

@@ -19,4 +19,5 @@ dependencies = [
     "whitenoise>=6.8.2",
     "gunicorn>=23.0.0",
     "django-stubs>=5.1.1",
+    "pdbpp>=0.10.3",
 ]