Ensure task compression actually happens when setting task_compression (#7470)

thedrow · Omer Katz · web-flow · commit ae20aa9d0668 · 2022-04-18T22:15:05.000+06:00
* Ensure task compression actually happens when setting `task_compression`. Fixes #4838. Previously, we erroneously used `result_compression` as the configuration option for this behavior. It appears that compressing results was never supported in Celery or that the support for it was removed. This will be fixed later on. * Happify the linter. Co-authored-by: Omer Katz <omer.katz@kcg.tech>
diff --git a/celery/app/amqp.py b/celery/app/amqp.py
@@ -447,7 +447,7 @@ def _create_task_sender(self):
 
         default_rkey = self.app.conf.task_default_routing_key
         default_serializer = self.app.conf.task_serializer
-        default_compressor = self.app.conf.result_compression
+        default_compressor = self.app.conf.task_compression
 
         def send_task_message(producer, name, message,
                               exchange=None, routing_key=None, queue=None,
diff --git a/t/unit/app/test_amqp.py b/t/unit/app/test_amqp.py
@@ -205,8 +205,7 @@ def test_as_task_message_without_utc(self):
         self.app.amqp.as_task_v1(uuid(), 'foo', countdown=30, expires=40)
 
 
-class test_AMQP:
-
+class test_AMQP_Base:
     def setup(self):
         self.simple_message = self.app.amqp.as_task_v2(
             uuid(), 'foo', create_sent_event=True,
@@ -215,6 +214,9 @@ def setup(self):
             uuid(), 'foo', create_sent_event=False,
         )
 
+
+class test_AMQP(test_AMQP_Base):
+
     def test_kwargs_must_be_mapping(self):
         with pytest.raises(TypeError):
             self.app.amqp.as_task_v2(uuid(), 'foo', kwargs=[1, 2])
@@ -336,7 +338,7 @@ def update_conf_runtime_for_tasks_queues(self):
         assert router != router_was
 
 
-class test_as_task_v2:
+class test_as_task_v2(test_AMQP_Base):
 
     def test_raises_if_args_is_not_tuple(self):
         with pytest.raises(TypeError):
@@ -368,8 +370,27 @@ def test_eta_to_datetime(self):
         )
         assert m.headers['eta'] == eta.isoformat()
 
-    def test_callbacks_errbacks_chord(self):
+    def test_compression(self):
+        self.app.conf.task_compression = 'gzip'
+
+        prod = Mock(name='producer')
+        self.app.amqp.send_task_message(
+            prod, 'foo', self.simple_message_no_sent_event,
+            compression=None
+        )
+        assert prod.publish.call_args[1]['compression'] == 'gzip'
+
+    def test_compression_override(self):
+        self.app.conf.task_compression = 'gzip'
+
+        prod = Mock(name='producer')
+        self.app.amqp.send_task_message(
+            prod, 'foo', self.simple_message_no_sent_event,
+            compression='bz2'
+        )
+        assert prod.publish.call_args[1]['compression'] == 'bz2'
 
+    def test_callbacks_errbacks_chord(self):
         @self.app.task
         def t(i):
             pass
