fix(celery): Un-break usages of apply_async (getsentry#357)

untitaker · web-flow · commit 74b4405da6b3 · 2019-04-29T19:47:51.000+02:00
diff --git a/sentry_sdk/integrations/celery.py b/sentry_sdk/integrations/celery.py
@@ -51,7 +51,7 @@ def sentry_build_tracer(name, task, *args, **kwargs):
 
 
 def _wrap_apply_async(task, f):
-    def apply_async(self, *args, **kwargs):
+    def apply_async(*args, **kwargs):
         hub = Hub.current
         integration = hub.get_integration(CeleryIntegration)
         if integration is not None and integration.propagate_traces:
@@ -62,7 +62,7 @@ def apply_async(self, *args, **kwargs):
                 headers[key] = value
             if headers is not None:
                 kwargs["headers"] = headers
-        return f(self, *args, **kwargs)
+        return f(*args, **kwargs)
 
     return apply_async
 
diff --git a/tests/integrations/celery/test_celery.py b/tests/integrations/celery/test_celery.py
@@ -40,7 +40,22 @@ def celery(init_celery):
     return init_celery()
 
 
-def test_simple(capture_events, celery):
+@pytest.mark.parametrize(
+    "invocation,expected_context",
+    [
+        [lambda task, x, y: task.delay(x, y), {"args": [1, 0], "kwargs": {}}],
+        [lambda task, x, y: task.apply_async((x, y)), {"args": [1, 0], "kwargs": {}}],
+        [
+            lambda task, x, y: task.apply_async(args=(x, y)),
+            {"args": [1, 0], "kwargs": {}},
+        ],
+        [
+            lambda task, x, y: task.apply_async(kwargs=dict(x=x, y=y)),
+            {"args": [], "kwargs": {"x": 1, "y": 0}},
+        ],
+    ],
+)
+def test_simple(capture_events, celery, invocation, expected_context):
     events = capture_events()
 
     @celery.task(name="dummy_task")
@@ -51,18 +66,17 @@ def dummy_task(x, y):
     span_context = SpanContext.start_trace()
     with configure_scope() as scope:
         scope.set_span_context(span_context)
-    dummy_task.delay(1, 2)
-    dummy_task.delay(1, 0)
+
+    invocation(dummy_task, 1, 2)
+    invocation(dummy_task, 1, 0)
 
     event, = events
     assert event["contexts"]["trace"]["trace_id"] == span_context.trace_id
     assert event["contexts"]["trace"]["span_id"] != span_context.span_id
     assert event["transaction"] == "dummy_task"
-    assert event["extra"]["celery-job"] == {
-        "args": [1, 0],
-        "kwargs": {},
-        "task_name": "dummy_task",
-    }
+    assert event["extra"]["celery-job"] == dict(
+        task_name="dummy_task", **expected_context
+    )
 
     exception, = event["exception"]["values"]
     assert exception["type"] == "ZeroDivisionError"
