From 28f4d7d56641fed60c3c8d9bdf1d1f9d35fc6d5a Mon Sep 17 00:00:00 2001 From: Marc 'risson' Schmitt Date: Fri, 28 Mar 2025 17:16:34 +0100 Subject: [PATCH] fix tests Signed-off-by: Marc 'risson' Schmitt --- authentik/admin/tests/test_tasks.py | 1 + authentik/root/test_runner.py | 4 ++++ authentik/tasks/apps.py | 29 +++++++++++++++++++++++------ authentik/tasks/broker.py | 7 +++++-- authentik/tasks/middleware.py | 3 ++- authentik/tasks/test.py | 7 +++---- 6 files changed, 38 insertions(+), 13 deletions(-) diff --git a/authentik/admin/tests/test_tasks.py b/authentik/admin/tests/test_tasks.py index 126e9fbb8d..fd8ab2ff48 100644 --- a/authentik/admin/tests/test_tasks.py +++ b/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 diff --git a/authentik/root/test_runner.py b/authentik/root/test_runner.py index 867469d2b0..5816e61e3a 100644 --- a/authentik/root/test_runner.py +++ b/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""" diff --git a/authentik/tasks/apps.py b/authentik/tasks/apps.py index 28849902e5..b53d29c45a 100644 --- a/authentik/tasks/apps.py +++ b/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) diff --git a/authentik/tasks/broker.py b/authentik/tasks/broker.py index bcbad433bc..e783de33dc 100644 --- a/authentik/tasks/broker.py +++ b/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() diff --git a/authentik/tasks/middleware.py b/authentik/tasks/middleware.py index 36a659fdc9..6d44520d76 100644 --- a/authentik/tasks/middleware.py +++ b/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) diff --git a/authentik/tasks/test.py b/authentik/tasks/test.py index e8100e03ea..86ab32668f 100644 --- a/authentik/tasks/test.py +++ b/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)