Skip to content

Commit

Permalink
Rename function _wrap_apply_async to _wrap_task_run
Browse files Browse the repository at this point in the history
  • Loading branch information
divaltor committed Jan 24, 2024
1 parent c6c36ff commit 016609c
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 12 deletions.
13 changes: 8 additions & 5 deletions sentry_sdk/integrations/celery.py
Original file line number Diff line number Diff line change
Expand Up @@ -110,8 +110,8 @@ def sentry_build_tracer(name, task, *args, **kwargs):
from celery.app.task import Task # type: ignore
from celery import Celery # type: ignore

Check warning on line 111 in sentry_sdk/integrations/celery.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/celery.py#L111

Added line #L111 was not covered by tests

Task.apply_async = _wrap_apply_async(Task.apply_async)
Celery.send_task = _wrap_apply_async(Celery.send_task)
Task.apply_async = _wrap_task_run(Task.apply_async)
Celery.send_task = _wrap_task_run(Celery.send_task)

Check warning on line 114 in sentry_sdk/integrations/celery.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/celery.py#L113-L114

Added lines #L113 - L114 were not covered by tests

_patch_worker_exit()

Expand Down Expand Up @@ -146,7 +146,7 @@ def __exit__(self, exc_type, exc_value, traceback):
return None


def _wrap_apply_async(f):
def _wrap_task_run(f):

Check warning on line 149 in sentry_sdk/integrations/celery.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/celery.py#L149

Added line #L149 was not covered by tests
# type: (F) -> F
@wraps(f)
def apply_async(*args, **kwargs):
Expand All @@ -172,10 +172,13 @@ def apply_async(*args, **kwargs):
except (IndexError, TypeError):
task_started_from_beat = False

task = args[0]
if isinstance(args[0], Task):
task_name = args[0].name

Check warning on line 176 in sentry_sdk/integrations/celery.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/celery.py#L176

Added line #L176 was not covered by tests
else:
task_name = args[1]

Check warning on line 178 in sentry_sdk/integrations/celery.py

View check run for this annotation

Codecov / codecov/patch

sentry_sdk/integrations/celery.py#L178

Added line #L178 was not covered by tests

span_mgr = (
hub.start_span(op=OP.QUEUE_SUBMIT_CELERY, description=task.name)
hub.start_span(op=OP.QUEUE_SUBMIT_CELERY, description=task_name)
if not task_started_from_beat
else NoOpMgr()
) # type: Union[Span, NoOpMgr]
Expand Down
14 changes: 7 additions & 7 deletions tests/integrations/celery/test_celery.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
from sentry_sdk.integrations.celery import (
CeleryIntegration,
_get_headers,
_wrap_apply_async,
_wrap_task_run,
)

from sentry_sdk._compat import text_type
Expand Down Expand Up @@ -34,7 +34,7 @@ def init_celery(sentry_init, request):
def inner(propagate_traces=True, backend="always_eager", **kwargs):
sentry_init(
integrations=[CeleryIntegration(propagate_traces=propagate_traces)],
**kwargs
**kwargs,
)
celery = Celery(__name__)

Expand Down Expand Up @@ -372,16 +372,16 @@ def test_redis_backend_trace_propagation(
runs = []

@celery.task(name="dummy_task", bind=True)
def dummy_task(self):
def dummy_task(self, x, y):
runs.append(1)
1 / 0

with start_transaction(name="submit_celery"):
# Curious: Cannot use delay() here or py2.7-celery-4.2 crashes
if execution_way == "apply_async":
res = dummy_task.apply_async()
res = dummy_task.apply_async(kwargs={"x": 1, "y": 0})
elif execution_way == "send_task":
res = celery.send_task("dummy_task")
res = celery.send_task("dummy_task", kwargs={"x": 1, "y": 0})
else: # pragma: no cover
raise ValueError(execution_way)

Expand Down Expand Up @@ -579,7 +579,7 @@ def dummy_function(*args, **kwargs):
assert "sentry-trace" in headers
assert "baggage" in headers

wrapped = _wrap_apply_async(dummy_function)
wrapped = _wrap_task_run(dummy_function)
wrapped(mock.MagicMock(), (), headers={})


Expand All @@ -593,7 +593,7 @@ def dummy_function(*args, **kwargs):
assert "sentry-trace" not in headers
assert "baggage" not in headers

wrapped = _wrap_apply_async(dummy_function)
wrapped = _wrap_task_run(dummy_function)
wrapped(
mock.MagicMock(),
[
Expand Down

0 comments on commit 016609c

Please sign in to comment.