Add support for Celery 5

selinon/config.py

@@ -12,9 +12,12 @@
 import runpy
 import tempfile
 
+import celery
+
 from .errors import ConfigNotInitializedError
 from .errors import ConfigurationError
 from .errors import UnknownStorageError
+from .errors import UnsupportedCeleryError
 from .system import System
 from .trace import Trace
 
@@ -211,8 +214,18 @@ def set_celery_app(cls, celery_app):
         cls._logger.debug("Registering Selinon to Celery context")
 
         cls.celery_app = celery_app
-        celery_app.tasks.register(Dispatcher())
-        celery_app.tasks.register(SelinonTaskEnvelope())
+
+        celery_major_version = int(celery.__version__.split(".", maxsplit=1)[0])
+        if celery_major_version == 4:
+            celery_app.tasks.register(Dispatcher())
+            celery_app.tasks.register(SelinonTaskEnvelope())
+        elif celery_major_version == 5:
+            celery_app.register_task(Dispatcher())
+            celery_app.register_task(SelinonTaskEnvelope())
+        else:
+            raise UnsupportedCeleryError(
+                "Unsupported Celery version {}, supported are celery>=4,<6".format(celery.__version__)
+            )
 
     @classmethod
     def init(cls, celery_app, nodes_definition_file, flow_definition_files, config_py=None, keep_config_py=False):

selinon/errors.py

@@ -13,6 +13,10 @@ class SelinonException(Exception):
     """Base Selinon exception in exception hierarchy."""
 
 
+class UnsupportedCeleryError(Exception):
+    """An exception raised when using unsupported Celery version."""
+
+
 class FatalTaskError(SelinonException):
     """An exception that is raised by task on fatal error - task will be not retried."""
 

setup.py

@@ -44,7 +44,7 @@ def get_long_description():
     license='BSD',
     keywords='selinon celery yaml flow distributed-computing',
     extras_require={
-        'celery': ['celery>=4'],
+        'celery': ['celery>=4,<6'],
         'mongodb': ['pymongo'],
         'postgresql': ['SQLAlchemy', 'SQLAlchemy-Utils'],
         'redis': ['redis'],

test/test_retry.py

@@ -6,7 +6,7 @@
 # ######################################################################
 
 import pytest
-import flexmock
+from flexmock import flexmock
 from selinon_test_case import SelinonTestCase
 from selinon import SystemState
 from selinon.errors import DispatcherRetry