Fix custom headers propagation for protocol 1 hybrid messages

celery · Sep 29, 2020 · 14e5a84 · 14e5a84
1 parent ea37db1
commit 14e5a84
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
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_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"