set tags from scope

Author: snopoke

taskbadger/mug.py

@@ -83,18 +83,23 @@ class Scope:
     def __init__(self):
         self.stack = []
         self.context = {}
+        self.tags = {}
 
     def __enter__(self):
-        self.stack.append(self.context)
+        self.stack.append((self.context, self.tags))
         self.context = self.context.copy()
+        self.tags = self.tags.copy()
         return self
 
     def __exit__(self, *args):
-        self.context = self.stack.pop()
+        self.context, self.tags = self.stack.pop()
 
     def __setitem__(self, key, value):
         self.context[key] = value
 
+    def tag(self, tags: dict[str, str]):
+        self.tags.update(tags)
+
 
 class MugMeta(type):
     @property

taskbadger/sdk.py

@@ -130,14 +130,15 @@ def create_task(
         max_runtime=max_runtime,
         stale_timeout=stale_timeout,
     )
-    scope_data = Badger.current.scope().context
-    if scope_data or data:
+    scope = Badger.current.scope()
+    if scope.context or data:
         data = data or {}
-        task.data = {**scope_data, **data}
+        task.data = {**scope.context, **data}
     if actions:
         task.additional_properties = {"actions": [a.to_dict() for a in actions]}
-    if tags:
-        task.tags = TaskRequestTags.from_dict(tags)
+    if scope.tags or tags:
+        tags = tags or {}
+        task.tags = TaskRequestTags.from_dict({**scope.tags, **tags})
     kwargs = _make_args(body=task)
     if monitor_id:
         kwargs["x_taskbadger_monitor"] = monitor_id

tests/test_celery.py

@@ -77,11 +77,13 @@ def add_with_task_args(self, a, b):
             (2, 2),
             taskbadger_name="new_name",
             taskbadger_value_max=10,
-            taskbadger_kwargs={"data": {"foo": "bar"}},
+            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"}, status=StatusEnum.PENDING)
+    create.assert_called_once_with(
+        "new_name", value_max=10, data={"foo": "bar"}, tags={"bar": "baz"}, status=StatusEnum.PENDING
+    )
 
 
 def test_celery_task_with_kwargs(celery_session_app, celery_session_worker, bind_settings):

tests/test_scope.py

@@ -9,46 +9,58 @@ def test_scope_singleton():
     assert Badger.current == GLOBAL_MUG
     scope = Badger.current.scope()
     assert scope.context == {}
+    assert scope.tags == {}
     assert scope.stack == []
     assert scope == Badger.current.scope()
 
 
 def test_scope_context():
     scope = Badger.current.scope()
     assert scope.context == {}
+    assert scope.tags == {}
     assert scope.stack == []
     with scope:
-        assert scope.stack == [{}]
+        assert scope.stack == [({}, {})]
         scope.context["foo"] = "bar"
+        scope.tags["name"] = "value"
         with scope:
-            assert scope.stack == [{}, {"foo": "bar"}]
+            assert scope.stack == [({}, {}), ({"foo": "bar"}, {"name": "value"})]
             assert scope.context == {"foo": "bar"}
+            assert scope.tags == {"name": "value"}
             scope.context["bar"] = "bazz"
+            scope.tags["bar"] = "bazz"
             with scope:
                 assert scope.context == {"foo": "bar", "bar": "bazz"}
+                assert scope.tags == {"name": "value", "bar": "bazz"}
                 scope.context.clear()
+                scope.tags.clear()
         assert scope.context == {"foo": "bar"}
-        assert scope.stack == [{}]
+        assert scope.tags == {"name": "value"}
+        assert scope.stack == [({}, {})]
     assert scope.context == {}
+    assert scope.tags == {}
     assert scope.stack == []
 
 
 @pytest.fixture(autouse=True)
-def init_skd():
+def _init_skd():
     init("org", "project", "token")
 
 
 def test_create_task_with_scope(httpx_mock):
     with Badger.current.scope() as scope:
         scope["foo"] = "bar"
         scope["bar"] = "bazz"
+        scope.tag({"name": "value"})
         httpx_mock.add_response(
             url="https://taskbadger.net/api/org/project/tasks/",
             method="POST",
             match_headers={"Authorization": "Bearer token"},
-            match_content=b'{"name": "name", "status": "pending", "data": {"foo": "bar", "bar": "buzzer"}}',
+            match_content=b'{"name": "name", "status": "pending", '
+            b'"data": {"foo": "bar", "bar": "buzzer"}, '
+            b'"tags": {"name": "value", "name1": "value1"}}',
             json=_json_task_response(),
             status_code=201,
         )
-        task = create_task("name", data={"bar": "buzzer"})
+        task = create_task("name", data={"bar": "buzzer"}, tags={"name1": "value1"})
         _verify_task(task)