Fix custom headers propagation for protocol 1 hybrid messages

celery · Sep 30, 2020 · 0f1a53b · 0f1a53b
1 parent cd8782f
commit 0f1a53b
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 3 deletions.
diff --git a/celery/worker/strategy.py b/celery/worker/strategy.py
@@ -50,6 +50,7 @@ def hybrid_to_proto2(message, body):
         'kwargsrepr': body.get('kwargsrepr'),
         'origin': body.get('origin'),
     }
+    headers.update(message.headers or {})
 
     embed = {
         'callbacks': body.get('callbacks'),

diff --git a/t/unit/worker/test_request.py b/t/unit/worker/test_request.py
@@ -1204,8 +1204,8 @@ def test_execute_using_pool_with_none_timelimit_header(self):
 
     def test_execute_using_pool__defaults_of_hybrid_to_proto2(self):
         weakref_ref = Mock(name='weakref.ref')
-        headers = strategy.hybrid_to_proto2('', {'id': uuid(),
-                                                 'task': self.mytask.name})[1]
+        headers = strategy.hybrid_to_proto2(Mock(headers=None), {'id': uuid(),
+                                            'task': self.mytask.name})[1]
         job = self.zRequest(revoked_tasks=set(), ref=weakref_ref, **headers)
         job.execute_using_pool(self.pool)
         assert job._apply_result

diff --git a/t/unit/worker/test_strategy.py b/t/unit/worker/test_strategy.py
@@ -271,7 +271,7 @@ def failed():
 class test_hybrid_to_proto2:
 
     def setup(self):
-        self.message = Mock(name='message')
+        self.message = Mock(name='message', headers={"custom": "header"})
         self.body = {
             'args': (1,),
             'kwargs': {'foo': 'baz'},
@@ -288,3 +288,7 @@ def test_retries_custom_value(self):
         self.body['retries'] = _custom_value
         _, headers, _, _ = hybrid_to_proto2(self.message, self.body)
         assert headers.get('retries') == _custom_value
+
+    def test_custom_headers(self):
+        _, headers, _, _ = hybrid_to_proto2(self.message, self.body)
+        assert headers.get("custom") == "header"