From 7a5bc2423a15137a63d0ddcc96c16b6956f40f5e Mon Sep 17 00:00:00 2001 From: Simon Kelly Date: Thu, 1 Aug 2024 12:54:33 +0200 Subject: [PATCH 1/2] don't put TB kwargs in header if TB not configured This fixes a serialization bug if Actions are included in the TB configuration since they are not JSON serializable --- taskbadger/celery.py | 9 ++++++--- tests/test_celery.py | 9 ++++++++- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/taskbadger/celery.py b/taskbadger/celery.py index 3705121..c741b58 100644 --- a/taskbadger/celery.py +++ b/taskbadger/celery.py @@ -102,7 +102,6 @@ def scrape_urls(self, urls): def apply_async(self, *args, **kwargs): headers = kwargs.setdefault("headers", {}) - headers["taskbadger_track"] = True tb_kwargs = self._get_tb_kwargs(kwargs) if kwargs.get("kwargs"): # extract taskbadger options from task kwargs when supplied as keyword argument @@ -110,7 +109,11 @@ def apply_async(self, *args, **kwargs): elif len(args) > 1 and isinstance(args[1], dict): # extract taskbadger options from task kwargs when supplied as positional argument tb_kwargs.update(self._get_tb_kwargs(args[1])) - headers[TB_KWARGS_ARG] = tb_kwargs + + if Badger.is_configured(): + headers["taskbadger_track"] = True + headers[TB_KWARGS_ARG] = tb_kwargs + result = super().apply_async(*args, **kwargs) tb_task_id = result.info.get(TB_TASK_ID) if result.info else None @@ -150,13 +153,13 @@ def taskbadger_task(self): @before_task_publish.connect def task_publish_handler(sender=None, headers=None, body=None, **kwargs): headers = headers if "task" in headers else body + header_kwargs = headers.pop(TB_KWARGS_ARG, {}) # always remove TB headers if sender.startswith("celery.") or not Badger.is_configured(): return celery_system = Badger.current.settings.get_system_by_id("celery") auto_track = celery_system and celery_system.track_task(sender) manual_track = headers.get("taskbadger_track") - header_kwargs = headers.pop(TB_KWARGS_ARG, {}) if not manual_track and not auto_track: return diff --git a/tests/test_celery.py b/tests/test_celery.py index dc1158c..cb50268 100644 --- a/tests/test_celery.py +++ b/tests/test_celery.py @@ -171,7 +171,14 @@ def add_no_tb(self, a, b): 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) + 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() From 483d96beccfe9fae6b82f67f8e017788505bef39 Mon Sep 17 00:00:00 2001 From: Simon Kelly Date: Thu, 1 Aug 2024 12:54:42 +0200 Subject: [PATCH 2/2] ignore test thread class in tests --- tests/test_session.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_session.py b/tests/test_session.py index da6cd3d..8ab7039 100644 --- a/tests/test_session.py +++ b/tests/test_session.py @@ -71,6 +71,8 @@ def test_session_multiple_threads(): class TestThread(threading.Thread): + __test__ = False + def __init__(self, name, barrier, clients): threading.Thread.__init__(self, name=name) self.barrier = barrier