test_celery.py

"""
Note
====

As part of the Celery fixture setup a 'ping' task is run which executes
before the `_bind_settings` fixture is executed. This means that if any code
calls `Badger.is_configured()` (or similar), the `_local` ContextVar in the
Celery runner thread will not have the configuration set.
"""

import logging
import time
from unittest import mock

import celery
import pytest
from kombu.utils.json import register_type

from taskbadger import Action, EmailIntegration, StatusEnum
from taskbadger.celery import Task, task_publish_handler
from taskbadger.mug import Badger, Settings
from taskbadger.systems.celery import CelerySystemIntegration
from tests.utils import task_for_test


def _wait_for_call_count(mock_obj, expected, timeout=10.0):
    """Wait until ``mock_obj`` has been called ``expected`` times.

    The worker's ``task_success`` signal fires asynchronously, so the final
    status update can lag behind ``result.get()`` returning. Poll for it while
    the mock patch is still active to avoid a race with the ``with`` block exit.
    """
    deadline = time.monotonic() + timeout
    while mock_obj.call_count < expected and time.monotonic() < deadline:
        time.sleep(0.01)


@pytest.fixture(autouse=True)
def _check_log_errors(caplog):
    yield
    errors = [r.getMessage() for r in caplog.get_records("call") if r.levelno == logging.ERROR]
    if errors:
        pytest.fail(f"log errors during tests: {errors}")


@pytest.mark.usefixtures("_bind_settings")
def test_celery_task(celery_session_app, celery_session_worker):
    @celery_session_app.task(bind=True, base=Task)
    def add_normal(self, a, b):
        assert self.request.get("taskbadger_task") is not None, "missing task in request"
        assert self.taskbadger_task is not None, "missing task on self"
        assert Badger.current.session().client is not None, "missing client"
        return a + b

    celery_session_worker.reload()

    with (
        mock.patch("taskbadger.celery.create_task_safe") as create,
        mock.patch("taskbadger.celery.update_task_safe") as update,
        mock.patch("taskbadger.sdk.get_task") as get_task,
    ):
        tb_task = task_for_test()
        create.return_value = tb_task
        result = add_normal.delay(2, 2)
        assert result.info.get("taskbadger_task_id") == tb_task.id
        assert result.taskbadger_task_id == tb_task.id
        assert result.get_taskbadger_task() is not None
        assert result.get(timeout=10, propagate=True) == 4
        _wait_for_call_count(update, 2)

    create.assert_called_once()
    assert get_task.call_count == 2
    assert update.call_count == 2
    assert Badger.current.session().client is None


@pytest.mark.usefixtures("_bind_settings")
def test_celery_task_with_args(celery_session_app, celery_session_worker):
    @celery_session_app.task(bind=True, base=Task)
    def add_with_task_args(self, a, b):
        assert self.taskbadger_task is not None
        return a + b

    celery_session_worker.reload()

    with (
        mock.patch("taskbadger.celery.create_task_safe") as create,
        mock.patch("taskbadger.celery.update_task_safe"),
        mock.patch("taskbadger.sdk.get_task"),
    ):
        create.return_value = task_for_test()

        result = add_with_task_args.apply_async(
            (2, 2),
            taskbadger_name="new_name",
            taskbadger_value_max=10,
            taskbadger_kwargs={"data": {"foo": "bar"}, "tags": {"bar": "baz"}},
        )
        assert result.get(timeout=10, propagate=True) == 4

    create.assert_called_once_with(
        "new_name",
        value_max=10,
        data={"foo": "bar"},
        tags={"bar": "baz"},
        status=StatusEnum.PENDING,
        queue="celery",
        external_id=mock.ANY,
    )


@pytest.mark.usefixtures("_bind_settings")
def test_celery_task_with_kwargs(celery_session_app, celery_session_worker):
    @celery_session_app.task(bind=True, base=Task)
    def add_with_task_args(self, a, b):
        assert self.taskbadger_task is not None
        return a + b

    celery_session_worker.reload()

    with (
        mock.patch("taskbadger.celery.create_task_safe") as create,
        mock.patch("taskbadger.celery.update_task_safe"),
        mock.patch("taskbadger.sdk.get_task"),
    ):
        create.return_value = task_for_test()

        actions = [Action("stale", integration=EmailIntegration(to="test@test.com"))]
        result = add_with_task_args.delay(
            2,
            2,
            taskbadger_name="new_name",
            taskbadger_value_max=10,
            taskbadger_kwargs={"actions": actions},
        )
        assert result.get(timeout=10, propagate=True) == 4

    create.assert_called_once_with(
        "new_name", value_max=10, actions=actions, status=StatusEnum.PENDING, queue="celery", external_id=mock.ANY
    )


@pytest.mark.usefixtures("_bind_settings")
def test_celery_record_args(celery_session_app, celery_session_worker):
    @celery_session_app.task(bind=True, base=Task)
    def add_with_task_args(self, a, b):
        assert self.taskbadger_task is not None
        return a + b

    celery_session_worker.reload()

    with (
        mock.patch("taskbadger.celery.create_task_safe") as create,
        mock.patch("taskbadger.celery.update_task_safe"),
        mock.patch("taskbadger.sdk.get_task"),
    ):
        create.return_value = task_for_test()

        result = add_with_task_args.apply_async(
            (2, 2),
            taskbadger_name="new_name",
            taskbadger_value_max=10,
            taskbadger_kwargs={"data": {"foo": "bar"}},
            taskbadger_record_task_args=True,
        )
        assert result.get(timeout=10, propagate=True) == 4

    create.assert_called_once_with(
        "new_name",
        value_max=10,
        data={"foo": "bar", "celery_task_args": [2, 2], "celery_task_kwargs": {}},
        status=StatusEnum.PENDING,
        queue="celery",
        external_id=mock.ANY,
    )


@pytest.mark.usefixtures("_bind_settings")
def test_celery_record_task_kwargs(celery_session_app, celery_session_worker):
    @celery_session_app.task(bind=True, base=Task)
    def add_with_task_kwargs(self, a, b, c=0):
        assert self.taskbadger_task is not None
        return a + b + c

    celery_session_worker.reload()

    with (
        mock.patch("taskbadger.celery.create_task_safe") as create,
        mock.patch("taskbadger.celery.update_task_safe"),
        mock.patch("taskbadger.sdk.get_task"),
    ):
        create.return_value = task_for_test()

        actions = [Action("stale", integration=EmailIntegration(to="test@test.com"))]
        result = add_with_task_kwargs.delay(
            2,
            2,
            c=3,
            taskbadger_name="new_name",
            taskbadger_value_max=10,
            taskbadger_kwargs={"actions": actions},
            taskbadger_record_task_args=True,
        )
        assert result.get(timeout=10, propagate=True) == 7

    create.assert_called_once_with(
        "new_name",
        value_max=10,
        data={"celery_task_args": [2, 2], "celery_task_kwargs": {"c": 3}},
        actions=actions,
        status=StatusEnum.PENDING,
        queue="celery",
        external_id=mock.ANY,
    )


@pytest.mark.usefixtures("_bind_settings")
def test_celery_record_task_args_custom_serialization(celery_session_app, celery_session_worker):
    class A:
        def __init__(self, a, b):
            self.a = a
            self.b = b

    register_type(A, "A", lambda o: [o.a, o.b], lambda o: A(*o))

    @celery_session_app.task(bind=True, base=Task)
    def add_task_custom_serialization(self, a):
        assert self.taskbadger_task is not None
        return a.a + a.b

    celery_session_worker.reload()

    with (
        mock.patch("taskbadger.celery.create_task_safe") as create,
        mock.patch("taskbadger.celery.update_task_safe"),
        mock.patch("taskbadger.sdk.get_task"),
    ):
        create.return_value = task_for_test()

        result = add_task_custom_serialization.delay(
            A(2, 2),
            taskbadger_record_task_args=True,
        )
        assert result.get(timeout=10, propagate=True) == 4

    create.assert_called_once_with(
        "tests.test_celery.add_task_custom_serialization",
        data={"celery_task_args": [{"__type__": "A", "__value__": [2, 2]}], "celery_task_kwargs": {}},
        status=StatusEnum.PENDING,
        queue="celery",
        external_id=mock.ANY,
    )


@pytest.mark.usefixtures("_bind_settings")
def test_celery_task_with_args_in_decorator(celery_session_app, celery_session_worker):
    @celery_session_app.task(
        bind=True,
        base=Task,
        taskbadger_value_max=10,
        taskbadger_kwargs={"monitor_id": "123"},
    )
    def add_with_task_args_in_decorator(self, a, b):
        assert self.taskbadger_task is not None
        return a + b

    celery_session_worker.reload()

    with (
        mock.patch("taskbadger.celery.create_task_safe") as create,
        mock.patch("taskbadger.celery.update_task_safe"),
        mock.patch("taskbadger.sdk.get_task"),
    ):
        create.return_value = task_for_test()

        result = add_with_task_args_in_decorator.delay(2, 2)
        assert result.get(timeout=10, propagate=True) == 4

    create.assert_called_once_with(
        mock.ANY, status=StatusEnum.PENDING, monitor_id="123", value_max=10, queue="celery", external_id=mock.ANY
    )


@pytest.mark.usefixtures("_bind_settings")
def test_celery_publish_handler_task_not_registered_locally():
    """before_task_publish can fire in a process that never imported the task
    being sent (e.g. a producer-only service), so `celery.current_app.tasks.get`
    returns None. The handler should still create the TaskBadger task without
    blowing up trying to call `.update_state()` on the missing task class."""

    with mock.patch("taskbadger.celery.create_task_safe") as create:
        create.return_value = task_for_test()

        headers = {"id": "abc123", "task": "unregistered.task", "taskbadger_track": True}
        task_publish_handler(sender="unregistered.task", headers=headers, body=[[], {}, {}])

    create.assert_called_once()
    assert headers["taskbadger_task_id"] == create.return_value.id


@pytest.mark.parametrize(("track_header", "expect_created"), [({}, True), ({"taskbadger_track": False}, False)])
def test_celery_publish_handler_opt_out_beats_auto_track(track_header, expect_created):
    """An explicit `taskbadger_track=False` header opts a single execution out, even
    when auto-tracking would otherwise pick the task up."""
    settings = Settings("https://taskbadger.net", "token", "org", "proj", systems={"celery": CelerySystemIntegration()})
    Badger.current.bind(settings)
    try:
        with mock.patch("taskbadger.celery.create_task_safe") as create:
            create.return_value = task_for_test()
            headers = {"id": "abc123", "task": "auto.tracked.task", **track_header}
            task_publish_handler(sender="auto.tracked.task", headers=headers, body=[[], {}, {}])
    finally:
        Badger.current.bind(None)

    assert create.called is expect_created


@pytest.mark.usefixtures("_bind_settings")
def test_celery_track_attr_not_passed_to_create(celery_session_app):
    """`taskbadger_track` is a tracking switch, not a task field, so it must never reach
    `create_task` — which would raise on the unexpected kwarg and lose the task."""

    @celery_session_app.task(base=Task, name="track.attr.task", taskbadger_track=False)
    def track_attr_task():
        return 1

    with mock.patch("taskbadger.celery.create_task_safe") as create:
        create.return_value = task_for_test()
        headers = {"id": "abc123", "task": "track.attr.task", "taskbadger_track": True}
        task_publish_handler(sender="track.attr.task", headers=headers, body=[[], {}, {}])

    assert "track" not in create.call_args.kwargs


@pytest.mark.usefixtures("_bind_settings")
def test_celery_task_opt_out(celery_session_app, celery_session_worker):
    """`headers={"taskbadger_track": False}` prevents tracking of a single execution."""

    @celery_session_app.task(bind=True, base=Task)
    def add_opt_out(self, a, b):
        assert self.taskbadger_task_id is None, "task should not be tracked"
        return a + b

    celery_session_worker.reload()

    with (
        mock.patch("taskbadger.celery.create_task_safe") as create,
        mock.patch("taskbadger.celery.update_task_safe") as update,
    ):
        result = add_opt_out.apply_async((2, 2), headers={"taskbadger_track": False})
        assert result.get(timeout=10, propagate=True) == 4

    create.assert_not_called()
    update.assert_not_called()


@pytest.mark.usefixtures("_bind_settings")
def test_celery_task_opt_out_kwarg(celery_session_app, celery_session_worker):
    """`taskbadger_track=False` also works as a `taskbadger_`-prefixed option, the way the
    other per-call options are passed, and never leaks into the create_task kwargs."""

    @celery_session_app.task(bind=True, base=Task)
    def add_opt_out_kwarg(self, a, b):
        assert self.taskbadger_task_id is None, "task should not be tracked"
        return a + b

    celery_session_worker.reload()

    with (
        mock.patch("taskbadger.celery.create_task_safe") as create,
        mock.patch("taskbadger.celery.update_task_safe") as update,
    ):
        result = add_opt_out_kwarg.apply_async((2, 2), taskbadger_track=False)
        assert result.get(timeout=10, propagate=True) == 4

    create.assert_not_called()
    update.assert_not_called()


@pytest.mark.usefixtures("_bind_settings")
def test_celery_task_custom_queue(celery_session_app, celery_session_worker):
    @celery_session_app.task(bind=True, base=Task)
    def add_custom_queue(self, a, b):
        return a + b

    celery_session_worker.reload()

    with (
        mock.patch("taskbadger.celery.create_task_safe") as create,
        mock.patch("taskbadger.celery.update_task_safe"),
        mock.patch("taskbadger.sdk.get_task"),
    ):
        create.return_value = task_for_test()
        # The task is routed to a queue the session worker doesn't consume, so we
        # only assert on the create call made synchronously during publish.
        add_custom_queue.apply_async((2, 2), queue="high_priority")

    assert create.call_args.kwargs["queue"] == "high_priority"


@pytest.mark.usefixtures("_bind_settings")
def test_celery_records_task_id_as_external_id(celery_session_app, celery_session_worker):
    @celery_session_app.task(bind=True, base=Task)
    def add_external_id(self, a, b):
        return a + b

    celery_session_worker.reload()

    with (
        mock.patch("taskbadger.celery.create_task_safe") as create,
        mock.patch("taskbadger.celery.update_task_safe"),
        mock.patch("taskbadger.sdk.get_task"),
    ):
        create.return_value = task_for_test()
        result = add_external_id.apply_async((2, 2), queue="high_priority")

    assert create.call_args.kwargs["external_id"] == result.id


@pytest.mark.usefixtures("_bind_settings")
def test_celery_external_id_not_cached_across_invocations(celery_session_app, celery_session_worker):
    # Tasks defined with taskbadger_kwargs share a class-level dict; external_id
    # must not leak from one publish to the next.
    @celery_session_app.task(bind=True, base=Task, taskbadger_kwargs={"value_max": 10})
    def add_repeat(self, a, b):
        return a + b

    celery_session_worker.reload()

    with (
        mock.patch("taskbadger.celery.create_task_safe") as create,
        mock.patch("taskbadger.celery.update_task_safe"),
        mock.patch("taskbadger.sdk.get_task"),
    ):
        create.return_value = task_for_test()
        result1 = add_repeat.apply_async((2, 2), queue="high_priority")
        result2 = add_repeat.apply_async((3, 3), queue="high_priority")

    external_ids = [c.kwargs["external_id"] for c in create.call_args_list]
    assert external_ids == [result1.id, result2.id]
    assert result1.id != result2.id


@pytest.mark.usefixtures("_bind_settings")
def test_celery_task_retry(celery_session_app, celery_session_worker):
    """Note: When a task is retried, the celery task ID remains the same but a new TB task
    will be created.

    TODO: How to handle this in TB? Should we update the existing TB task or create a new one?
    """

    @celery_session_app.task(bind=True, base=Task)
    def add_retry(self, a, b, is_retry=False):
        assert self.taskbadger_task is not None
        assert Badger.current.session().client is not None
        if not is_retry:
            raise self.retry(kwargs={"is_retry": True}, max_retries=1, countdown=0)

        return a + b

    celery_session_worker.reload()

    with (
        mock.patch("taskbadger.celery.create_task_safe") as create,
        mock.patch("taskbadger.celery.update_task_safe") as update,
        mock.patch("taskbadger.sdk.get_task") as get_task,
    ):
        create.return_value = task_for_test()
        get_task.return_value = task_for_test()
        result = add_retry.delay(2, 2)
        assert result.get(timeout=10, propagate=True) == 4
        _wait_for_call_count(update, 4)

    assert create.call_count == 2
    assert update.call_args_list == [
        mock.call(mock.ANY, status=StatusEnum.PROCESSING, data=mock.ANY),
        mock.call(mock.ANY, status=StatusEnum.ERROR, data=mock.ANY),
        mock.call(mock.ANY, status=StatusEnum.PROCESSING, data=mock.ANY),
        mock.call(mock.ANY, status=StatusEnum.SUCCESS, data=mock.ANY),
    ]
    assert Badger.current.session().client is None


def test_celery_task_badger_not_configured(celery_session_app, celery_session_worker):
    @celery_session_app.task(bind=True, base=Task)
    def add_no_tb(self, a, b):
        with mock.patch("taskbadger.sdk.get_task") as get_task:
            assert self.taskbadger_task_id is None
            assert Badger.current.session().client is None
            get_task.assert_not_called()

        return a + b

    celery_session_worker.reload()

    with (
        mock.patch("taskbadger.celery.create_task_safe") as create,
        mock.patch("taskbadger.celery.update_task_safe") as update,
    ):
        result = add_no_tb.delay(
            2,
            2,
            taskbadger_kwargs={
                # add an action here to test serialization failure when Badger is not configured
                "actions": [Action("stale", integration=EmailIntegration(to="test@test.com"))]
            },
        )
        assert result.get(timeout=10, propagate=True) == 4

    create.assert_not_called()
    update.assert_not_called()
    assert Badger.current.session().client is None


def test_task_direct_call(celery_session_app, celery_session_worker):
    @celery_session_app.task(bind=True, base=Task)
    def add_direct(self, a, b):
        with mock.patch("taskbadger.sdk.get_task") as get_task:
            assert self.taskbadger_task_id is None
            assert Badger.current.session().client is None
            get_task.assert_not_called()

        return a + b

    celery_session_worker.reload()

    assert add_direct(2, 1) == 3
    result = add_direct.apply(args=[2, 3])
    assert result.get(timeout=1, propagate=True) == 5
    assert Badger.current.session().client is None


@pytest.mark.usefixtures("_bind_settings")
def test_task_shared_task(celery_session_worker):
    @celery.shared_task(bind=True, base=Task)
    def add_shared_task(self, a, b):
        assert self.taskbadger_task is not None
        assert Badger.current.session().client is not None
        return a + b

    celery_session_worker.reload()

    with (
        mock.patch("taskbadger.celery.create_task_safe") as create,
        mock.patch("taskbadger.celery.update_task_safe") as update,
        mock.patch("taskbadger.sdk.get_task"),
    ):
        create.return_value = task_for_test()

        result = add_shared_task.delay(2, 2)
        assert result.get(timeout=10, propagate=True) == 4

    create.assert_called()
    update.assert_called()
    assert Badger.current.session().client is None


@pytest.mark.usefixtures("_bind_settings")
def test_task_signature(celery_session_worker):
    @celery.shared_task(bind=True, base=Task)
    def task_signature(self, a):
        assert self.taskbadger_task is not None
        assert Badger.current.session().client is not None
        return a * 2

    celery_session_worker.reload()

    chain = task_signature.s(2) | task_signature.s() | task_signature.s()

    with (
        mock.patch("taskbadger.celery.create_task_safe") as create,
        mock.patch("taskbadger.celery.update_task_safe") as update,
        mock.patch("taskbadger.sdk.get_task") as get_task,
    ):
        create.return_value = task_for_test()

        result = chain.delay()
        assert result.get(timeout=10, propagate=True) == 16
        _wait_for_call_count(update, 6)

    assert create.call_count == 3
    assert get_task.call_count == 3
    assert update.call_count == 6
    assert Badger.current.session().client is None


@pytest.mark.usefixtures("_bind_settings")
def test_task_map(celery_session_worker):
    """Tasks executed via map canvas primitive should be tracked.

    Note: The individual function calls within map are not separate Celery tasks,
    so we track the map operation itself with the inner task's name.
    """

    @celery.shared_task(bind=True, base=Task)
    def task_map_fn(self, a):
        return a * 2

    celery_session_worker.reload()

    map_canvas = task_map_fn.map(list(range(5)))

    with (
        mock.patch("taskbadger.celery.create_task_safe") as create,
        mock.patch("taskbadger.celery.update_task_safe") as update,
        mock.patch("taskbadger.sdk.get_task"),
    ):
        tb_task = task_for_test()
        create.return_value = tb_task
        result = map_canvas.delay()
        assert result.get(timeout=10, propagate=True) == [0, 2, 4, 6, 8]
        _wait_for_call_count(update, 2)

    # Map operation should create one TaskBadger task
    assert create.call_count == 1
    # Verify the task name includes inner task name and canvas type/count suffix
    call_args = create.call_args
    task_name = call_args[0][0]
    assert "task_map_fn" in task_name
    assert task_name.endswith("(map 5)")
    assert update.call_count == 2  # PROCESSING and SUCCESS
    assert Badger.current.session().client is None


@pytest.mark.usefixtures("_bind_settings")
def test_task_map_opt_out(celery_session_worker):
    """Canvas tasks honour the opt-out too. They are created in the worker rather than
    at publish time, so the header is checked there."""

    @celery.shared_task(bind=True, base=Task)
    def task_map_opt_out_fn(self, a):
        return a * 2

    celery_session_worker.reload()

    map_canvas = task_map_opt_out_fn.map(list(range(3)))

    with (
        mock.patch("taskbadger.celery.create_task_safe") as create,
        mock.patch("taskbadger.celery.update_task_safe") as update,
    ):
        result = map_canvas.apply_async(headers={"taskbadger_track": False})
        assert result.get(timeout=10, propagate=True) == [0, 2, 4]

    create.assert_not_called()
    update.assert_not_called()


@pytest.mark.usefixtures("_bind_settings")
def test_task_starmap(celery_session_worker):
    """Tasks executed via starmap canvas primitive should be tracked."""

    @celery.shared_task(bind=True, base=Task)
    def task_starmap_fn(self, a, b):
        return a + b

    celery_session_worker.reload()

    starmap_canvas = task_starmap_fn.starmap([(1, 2), (3, 4), (5, 6)])

    with (
        mock.patch("taskbadger.celery.create_task_safe") as create,
        mock.patch("taskbadger.celery.update_task_safe") as update,
        mock.patch("taskbadger.sdk.get_task"),
    ):
        tb_task = task_for_test()
        create.return_value = tb_task
        result = starmap_canvas.delay()
        assert result.get(timeout=10, propagate=True) == [3, 7, 11]
        _wait_for_call_count(update, 2)

    # Starmap operation should create one TaskBadger task
    assert create.call_count == 1
    # Verify the task name includes inner task name and canvas type/count suffix
    call_args = create.call_args
    task_name = call_args[0][0]
    assert "task_starmap_fn" in task_name
    assert task_name.endswith("(starmap 3)")
    assert update.call_count == 2  # PROCESSING and SUCCESS


@pytest.mark.usefixtures("_bind_settings")
def test_task_chunks(celery_session_worker):
    """Tasks executed via chunks canvas primitive should be tracked."""

    @celery.shared_task(bind=True, base=Task)
    def task_chunks_fn(self, a):
        return a * 2

    celery_session_worker.reload()

    # chunks creates multiple starmap tasks
    chunks_canvas = task_chunks_fn.chunks([(x,) for x in range(6)], 2).group()

    with (
        mock.patch("taskbadger.celery.create_task_safe") as create,
        mock.patch("taskbadger.celery.update_task_safe") as update,
        mock.patch("taskbadger.sdk.get_task"),
    ):
        tb_task = task_for_test()
        create.return_value = tb_task
        result = chunks_canvas.delay()
        assert result.get(timeout=10, propagate=True) == [[0, 2], [4, 6], [8, 10]]
        _wait_for_call_count(update, 6)

    # Each chunk should create a TaskBadger task (3 chunks of 2)
    assert create.call_count == 3
    # Verify the task names include inner task name and canvas type/count suffix
    for call in create.call_args_list:
        task_name = call[0][0]
        assert "task_chunks_fn" in task_name
        assert task_name.endswith("(starmap 2)")  # chunks uses starmap internally
    assert update.call_count == 6  # 3 tasks * 2 updates each


@pytest.mark.usefixtures("_bind_settings")
def test_celery_task_already_in_terminal_state(celery_session_worker):
    @celery.shared_task(bind=True, base=Task)
    def add_manual_update(self, a, b, is_retry=False):
        # simulate updating the task to a terminal state
        self.request.update({"taskbadger_task": task_for_test(status=StatusEnum.SUCCESS)})
        return a + b

    celery_session_worker.reload()

    with (
        mock.patch("taskbadger.celery.create_task_safe") as create,
        mock.patch("taskbadger.celery.update_task_safe") as update,
        mock.patch("taskbadger.sdk.get_task") as get_task,
    ):
        create.return_value = task_for_test()

        get_task.return_value = task_for_test()
        result = add_manual_update.delay(2, 2)
        assert result.get(timeout=10, propagate=True) == 4

    assert create.call_count == 1
    assert update.call_args_list == [
        mock.call(mock.ANY, status=StatusEnum.PROCESSING, data=mock.ANY),
    ]