fix tests

Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space>

authentik/admin/tests/test_tasks.py

@@ -1,5 +1,6 @@
 """test admin tasks"""
 
+import dramatiq
 from django.core.cache import cache
 from django.test import TestCase
 from requests_mock import Mocker

authentik/root/test_runner.py

@@ -4,6 +4,8 @@ import os
 from argparse import ArgumentParser
 from unittest import TestCase
 
+from django.apps import apps
+from django.test.testcases import apps
 import pytest
 from django.conf import settings
 from django.test.runner import DiscoverRunner
@@ -52,6 +54,8 @@ class PytestTestRunner(DiscoverRunner):  # pragma: no cover
         startup.send(sender=self, mode="test")
         post_startup.send(sender=self, mode="test")
 
+        apps.get_app_config("authentik_tasks").use_test_broker()
+
     @classmethod
     def add_arguments(cls, parser: ArgumentParser):
         """Add more pytest-specific arguments"""

authentik/tasks/apps.py

@@ -1,6 +1,6 @@
 import dramatiq
+from dramatiq.broker import Broker, get_broker
 from dramatiq.encoder import PickleEncoder
-from django.conf import settings
 
 from authentik.blueprints.apps import ManagedAppConfig
 
@@ -11,15 +11,32 @@ class AuthentikTasksConfig(ManagedAppConfig):
     verbose_name = "authentik Tasks"
     default = True
 
-    def ready(self) -> None:
-        from authentik.tasks.broker import PostgresBroker
-        from authentik.tasks.test import TestBroker
+    def _set_dramatiq_middlewares(self, broker: Broker) -> None:
         from authentik.tasks.middleware import CurrentTask, FullyQualifiedActorName
 
-        dramatiq.set_encoder(PickleEncoder())
-        broker = PostgresBroker() if not settings.TEST else TestBroker()
         broker.add_middleware(FullyQualifiedActorName())
         # broker.add_middleware(Prometheus())
         broker.add_middleware(CurrentTask())
+
+    def ready(self) -> None:
+        from authentik.tasks.broker import PostgresBroker
+
+        dramatiq.set_encoder(PickleEncoder())
+
+        broker = PostgresBroker()
+        self._set_dramatiq_middlewares(broker)
         dramatiq.set_broker(broker)
+
         return super().ready()
+
+    def use_test_broker(self) -> None:
+        from authentik.tasks.test import TestBroker
+
+        old_broker = get_broker()
+        broker = TestBroker()
+        self._set_dramatiq_middlewares(broker)
+        dramatiq.set_broker(broker)
+        for actor_name in old_broker.get_declared_actors():
+            actor = old_broker.get_actor(actor_name)
+            actor.broker = broker
+            actor.broker.declare_actor(actor)

authentik/tasks/broker.py

@@ -68,8 +68,11 @@ class DbConnectionMiddleware(Middleware):
     def _close_old_connections(self, *args, **kwargs):
         close_old_connections()
 
-    before_process_message = _close_old_connections
-    after_process_message = _close_old_connections
+    # TODO: figure out if we really need this, it seems a bit excessive to close connections after each message
+    # and if fails in tests
+
+    # before_process_message = _close_old_connections
+    # after_process_message = _close_old_connections
 
     def _close_connections(self, *args, **kwargs):
         connections.close_all()

authentik/tasks/middleware.py

@@ -41,5 +41,6 @@ class CurrentTask(Middleware):
         task: Task | None = self._TASK.get()
         if task is None:
             LOGGER.warn("Task was None, not saving. This should not happen")
-        task.save()
+        else:
+            task.save()
         self._TASK.set(None)

authentik/tasks/test.py

@@ -14,7 +14,7 @@ class TestWorker(Worker):
                 broker=self.broker,
                 queue_name=queue_name,
                 prefetch=2,
-                work_queue=None,
+                work_queue=self.work_queue,
                 worker_timeout=1,
             ),
         }
@@ -26,18 +26,17 @@ class TestWorker(Worker):
         self._worker = _WorkerThread(
             broker=self.broker,
             consumers=self.consumers,
-            work_queue=None,
+            work_queue=self.work_queue,
             worker_timeout=1,
         )
 
         self.broker.emit_before("worker_boot", self)
         self.broker.emit_after("worker_boot", self)
-        self.broker.emit_before("worker_thread_boot", self)
-        self.broker.emit_after("worker_thread_boot", self)
 
     def process_message(self, message: MessageProxy):
         self.logger.error(f"processing message {message}")
         self.work_queue.put(message)
+        self.consumers[message.queue_name].consumer.in_processing.add(message.message_id)
         self._worker.process_message(message)