Allow add_middleware to be called on brokers before init_app

CHANGES.rst

@@ -1,6 +1,12 @@
 Changelog
 =========
 
+Version 0.3.6
+-------------
+
+- Allow `add_middleware` to be called on brokers before `init_app`.
+
+
 Version 0.3.5
 -------------
 

docs/conf.py

@@ -62,9 +62,9 @@
 # built documents.
 #
 # The short X.Y version.
-version = '0.3.5'
+version = '0.3.6'
 # The full version, including alpha/beta/rc tags.
-release = '0.3.5'
+release = '0.3.6'
 
 # The language for content autogenerated by Sphinx. Refer to documentation
 # for a list of supported languages.
@@ -138,7 +138,7 @@
 # The name for this set of Sphinx documents.
 # "<project> v<release> documentation" by default.
 #
-# html_title = 'Flask-Melodramatiq v0.3.5'
+# html_title = 'Flask-Melodramatiq v0.3.6'
 
 # A shorter title for the navigation bar.  Default is the same as html_title.
 #

flask_melodramatiq/lazy_broker.py

@@ -108,11 +108,10 @@ def __init__(self, app=None, config_prefix=DEFAULT_CONFIG_PREFIX, **options):
         self.__options = options
         self.__configuration = None
 
-        # When an actor is defined, broker's `actor_options` attribute
-        # is accessed, which asks the middleware what actor options
-        # are valid. We will delegate this work to a stub broker,
-        # until our broker is ready.
-        self.__stub = dramatiq.brokers.stub.StubBroker(middleware=options.get('middleware'))
+        # We want to be able to add middleware and declare actors
+        # before `init_app` is called. We do this by delegating to a
+        # stub broker, until our broker is ready.
+        self.__stub = dramatiq.brokers.stub.StubBroker(middleware=options.pop('middleware', None))
 
         self._unregistered_lazy_actors = []
         if app is not None:
@@ -128,8 +127,9 @@ def init_app(self, app):
 
         """
 
-        configuration = self.__get_configuration(app)
         if self.__stub:
+            self.__options['middleware'] = self.__stub.middleware
+            configuration = self.__get_configuration(app)
             self.__stub.close()
             self.__stub = None
             self.__app = app
@@ -142,6 +142,8 @@ def init_app(self, app):
                 actor._register_proxied_instance(broker=broker)
             self._unregistered_lazy_actors = None
             self._proxied_instance = broker  # `self` is sealed from now on.
+        else:
+            configuration = self.__get_configuration(app)
         if configuration != self.__configuration:
             raise RuntimeError(
                 '{} tried to reconfigure an already configured broker.'.format(app)
@@ -192,6 +194,9 @@ def task():
     def actor_options(self):
         return (self._proxied_instance or self.__stub).actor_options
 
+    def add_middleware(self, middleware, *, before=None, after=None):
+        return (self._proxied_instance or self.__stub).add_middleware(middleware, before=before, after=after)
+
     def __get_primary_options(self):
         options = self.__options.copy()
         options.pop('class', None)

setup.py

@@ -22,7 +22,7 @@ def rel(*xs):
 
 setup(
     name='Flask-Melodramatiq',
-    version='0.3.5',
+    version='0.3.6',
     url='https://github.com/epandurski/flask_melodramatiq',
     license='MIT',
     author='Evgeni Pandurski',

tests/test_basics.py

@@ -3,6 +3,7 @@
 import pytest
 from flask_melodramatiq import Broker, StubBroker
 from flask_melodramatiq.lazy_broker import LazyActor, MultipleAppsWarningMiddleware
+from dramatiq.middleware import Middleware
 
 
 def test_actor_attr_access(app, broker, run_mock):
@@ -250,3 +251,28 @@ def test_after_process_boot_warning(broker, caplog):
     MultipleAppsWarningMiddleware().after_process_boot(broker)
     assert len(caplog.records) == n + 1
     assert 'application context may be set incorrectly' in caplog.text
+
+
+def test_add_middleware(app, broker, run_mock):
+    class TestMiddleware(Middleware):
+        @property
+        def actor_options(self):
+            return {"test_option"}
+
+        def after_process_message(self, broker, message, *, result=None, exception=None):
+            run_mock()
+
+    broker.add_middleware(TestMiddleware())
+
+    @broker.actor(test_option=123)
+    def task():
+        pass
+
+    broker.init_app(app)
+    broker.add_middleware(TestMiddleware())
+    assert task.options['test_option'] == 123
+    task.send()
+    worker = dramatiq.Worker(broker)
+    worker.start()
+    worker.join()
+    assert run_mock.call_count == 2