Add function to be used as a celery before_task_publish handler

Author: geniass

flask_log_request_id/extras/celery.py

@@ -24,6 +24,9 @@ def apply_async(self, *args, **kwargs):
 
         return super(RequestIDAwareTask, self).apply_async(*args, **kwargs)
 
+def add_request_id_header(headers=None, **kwargs):
+    if _CELERY_X_HEADER not in headers:
+        headers[_CELERY_X_HEADER] = current_request_id()
 
 def ctx_celery_task_get_request_id():
     """

tests/extras/celery_tests.py

@@ -1,10 +1,11 @@
-from celery import Celery
-import unittest
 import mock
+import unittest
 
-
-from flask_log_request_id.extras.celery import (
-    RequestIDAwareTask, ctx_celery_task_get_request_id, ExecutedOutsideContext)
+from celery import Celery
+from flask_log_request_id.extras.celery import (ExecutedOutsideContext,
+                                                RequestIDAwareTask,
+                                                add_request_id_header,
+                                                ctx_celery_task_get_request_id)
 
 
 class MockedTask(object):
@@ -63,6 +64,17 @@ def test_issue21_called_with_headers_None(self, mocked_current_request_id):
                     'foo': 'bar'
             })
 
+    @mock.patch('flask_log_request_id.extras.celery.current_request_id')
+    def test_before_task_publish_hooks_adds_header(self, mocked_current_request_id):
+        mocked_current_request_id.return_value = 15
+
+        headers = {}
+        add_request_id_header(headers={})
+        print(headers)
+        self.assertDictEqual(headers, {
+            'x_request_id': 15
+        })
+
     @mock.patch('flask_log_request_id.extras.celery.current_task')
     def test_ctx_fetcher_outside_context(self, mocked_current_task):