[Stress] Support both 3.1 and 4.0

funtests/stress/stress/app.py

@@ -1,13 +1,13 @@
 # -*- coding: utf-8 -*-
 from __future__ import absolute_import, print_function, unicode_literals
 
+import celery
 import os
 import sys
 import signal
 
 from time import sleep
 
-from celery import Celery
 from celery import signals
 from celery.bin.base import Option
 from celery.exceptions import SoftTimeLimitExceeded
@@ -17,8 +17,10 @@ from .templates import use_template, template_names
 
 logger = get_task_logger(__name__)
 
+IS_CELERY_4 = celery.VERSION[0] >= 4
 
-class App(Celery):
+
+class App(celery.Celery):
     template_selected = False
 
     def __init__(self, *args, **kwargs):
@@ -33,7 +35,8 @@ class App(Celery):
             )
         )
         signals.user_preload_options.connect(self.on_preload_parsed)
-        self.on_configure.connect(self._maybe_use_default_template)
+        if IS_CELERY_4:
+            self.on_configure.connect(self._maybe_use_default_template)
 
     def on_preload_parsed(self, options=None, **kwargs):
         self.use_template(options['template'])
@@ -48,6 +51,18 @@ class App(Celery):
         if not self.template_selected:
             self.use_template('default')
 
+    if not IS_CELERY_4:
+        after_configure = None
+
+        def _get_config(self):
+            ret = super(App, self)._get_config()
+            if self.after_configure:
+                self.after_configure(ret)
+            return ret
+
+        def on_configure(self):
+            self._maybe_use_default_template()
+
 app = App('stress', set_as_current=False)
 
 

funtests/stress/stress/templates.py

@@ -1,9 +1,12 @@
 from __future__ import absolute_import
 
+import celery
 import os
 
+from functools import partial
+
 from celery.five import items
-from kombu import Exchange, Queue
+from kombu import Queue
 from kombu.utils import symbol_by_name
 
 CSTRESS_TRANS = os.environ.get('CSTRESS_TRANS', False)
@@ -12,6 +15,8 @@ CSTRESS_QUEUE = os.environ.get('CSTRESS_QUEUE_NAME', default_queue)
 
 templates = {}
 
+IS_CELERY_4 = celery.VERSION[0] >= 4
+
 
 def template(name=None):
 
@@ -21,15 +26,23 @@ def template(name=None):
     return _register
 
 
-def use_template(app, template='default'):
-    template = template.split(',')
+if IS_CELERY_4:
+
+    def use_template(app, template='default'):
+        template = template.split(',')
+
+        # mixin the rest of the templates when the config is needed
+        @app.on_after_configure.connect(weak=False)
+        def load_template(sender, source, **kwargs):
+            mixin_templates(template[1:], source)
 
-    # mixin the rest of the templates when the config is needed
-    @app.on_after_configure.connect(weak=False)
-    def load_template(sender, source, **kwargs):
-        mixin_templates(template[1:], source)
+        app.config_from_object(templates[template[0]])
+else:
 
-    app.config_from_object(templates[template[0]])
+    def use_template(app, template='default'):  # noqa
+        template = template.split(',')
+        app.after_configure = partial(mixin_templates, template[1:])
+        app.config_from_object(templates[template[0]])
 
 
 def mixin_templates(templates, conf):