tests run locally

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

authentik/blueprints/v1/tasks.py

@@ -7,6 +7,7 @@ from sys import platform
 from uuid import UUID
 
 from dacite.core import from_dict
+from django.conf import settings
 from django.db import DatabaseError, InternalError, ProgrammingError
 from django.utils.text import slugify
 from django.utils.timezone import now
@@ -68,7 +69,8 @@ class BlueprintWatcherMiddleware(Middleware):
         observer.start()
 
     def after_worker_boot(self, broker, worker):
-        self.start_blueprint_watcher()
+        if not settings.TEST:
+            self.start_blueprint_watcher()
 
 
 class BlueprintEventHandler(FileSystemEventHandler):

authentik/root/test_runner.py

@@ -17,6 +17,7 @@ from authentik.events.context_processors.geoip import GEOIP_CONTEXT_PROCESSOR
 from authentik.lib.config import CONFIG
 from authentik.lib.sentry import sentry_init
 from authentik.root.signals import post_startup, pre_startup, startup
+from authentik.tasks.test import use_test_broker
 
 # globally set maxDiff to none to show full assert error
 TestCase.maxDiff = None
@@ -61,6 +62,7 @@ class PytestTestRunner(DiscoverRunner):  # pragma: no cover
     def _setup_test_environment(self):
         """Configure test environment settings"""
         settings.TEST = True
+        settings.DRAMATIQ["test"] = True
 
         # Test-specific configuration
         test_config = {
@@ -84,7 +86,7 @@ class PytestTestRunner(DiscoverRunner):  # pragma: no cover
         sentry_init()
         self.logger.debug("Test environment configured")
 
-        apps.get_app_config("authentik_tasks").use_test_broker()
+        use_test_broker()
 
         # Send startup signals
         pre_startup.send(sender=self, mode="test")

authentik/tasks/apps.py

@@ -9,18 +9,6 @@ class AuthentikTasksConfig(ManagedAppConfig):
     verbose_name = "authentik Tasks"
     default = True
 
-    # def use_test_broker(self) -> None:
-    #     from authentik.tasks.test import TestBroker
-    #
-    #     old_broker = get_broker()
-    #     broker = TestBroker(middleware=[])
-    #     self._set_dramatiq_middlewares(broker, max_retries=0)
-    #     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)
-
     @property
     def global_schedule_specs(self) -> list[ScheduleSpec]:
         from authentik.tasks.tasks import clean_worker_statuses

authentik/tasks/middleware.py

@@ -1,5 +1,4 @@
 import os
-from os import getpid
 from pathlib import Path
 from signal import pause
 from socket import gethostname
@@ -9,6 +8,7 @@ from typing import Any
 
 import pglock
 from django.utils.timezone import now
+from django.conf import settings
 from dramatiq.broker import Broker
 from dramatiq.common import current_millis
 from dramatiq.message import Message
@@ -186,6 +186,9 @@ class MetricsMiddleware(Middleware):
         return [worker_metrics]
 
     def before_worker_boot(self, broker: Broker, worker):
+        if settings.TEST:
+            return
+
         from prometheus_client import Counter, Gauge, Histogram
 
         self.total_messages = Counter(
@@ -252,7 +255,7 @@ class MetricsMiddleware(Middleware):
         from prometheus_client import multiprocess
 
         # TODO: worker_id
-        multiprocess.mark_process_dead(getpid())
+        multiprocess.mark_process_dead(os.getpid())
 
     def _make_labels(self, message: Message) -> tuple[str, str]:
         return (message.queue_name, message.actor_name)

authentik/tasks/test.py

@@ -1,9 +1,15 @@
 from queue import PriorityQueue
 
-from dramatiq.broker import Broker, MessageProxy
+from django.utils.module_loading import import_string
+from django_dramatiq_postgres.conf import Conf
+import dramatiq
+from dramatiq.broker import Broker, MessageProxy, get_broker
+from dramatiq.middleware.retries import Retries
+from dramatiq.results.middleware import Results
 from dramatiq.worker import Worker, _ConsumerThread, _WorkerThread
 
 from authentik.tasks.broker import PostgresBroker
+from authentik.tasks.middleware import MetricsMiddleware
 
 
 class TestWorker(Worker):
@@ -46,3 +52,31 @@ class TestBroker(PostgresBroker):
         worker = TestWorker(message.queue_name, broker=self)
         worker.process_message(MessageProxy(message))
         return message
+
+
+def use_test_broker():
+    old_broker = get_broker()
+
+    broker = TestBroker()
+
+    for actor_name in old_broker.get_declared_actors():
+        actor = old_broker.get_actor(actor_name)
+        actor.broker = broker
+        actor.broker.declare_actor(actor)
+
+    for middleware_class, middleware_kwargs in Conf().middlewares:
+        middleware: dramatiq.middleware.middleware.Middleware = import_string(middleware_class)(
+            **middleware_kwargs,
+        )
+        if isinstance(middleware, MetricsMiddleware):
+            continue
+        if isinstance(middleware, Retries):
+            middleware.max_retries = 0
+        if isinstance(middleware, Results):
+            middleware.backend = import_string(Conf().result_backend)(
+                *Conf().result_backend_args,
+                **Conf().result_backend_kwargs,
+            )
+        broker.add_middleware(middleware)
+
+    dramatiq.set_broker(broker)