Merge branch 'main' into donwillems/bump-to-0.1.5

kolonialno · Apr 20, 2023 · bc63449 · bc63449
2 parents 865917a + a503605
commit bc63449
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 2 deletions.
diff --git a/oida/commands/componentize.py b/oida/commands/componentize.py
@@ -239,14 +239,21 @@ def update_decorator(
     def leave_FunctionDef(
         self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef
     ) -> BaseStatement | FlattenSentinel[BaseStatement] | RemovalSentinel:
-        celery_task_decorator = m.Decorator(
+
+        # Matching tasks defined with @app.task
+        celery_app_task_decorator = m.Decorator(
             m.Call(m.Attribute(value=m.Name("app"), attr=m.Name("task")))
         )
+        # Matching tasks defined with @coalesced_task
+        celery_coalesced_task_decorator = m.Decorator(m.Call(m.Name("coalesced_task")))
+
         # The implicit name of the task is the path to the old module concatenated with the function name.
         task_name = f'"{self.module_name}.{original_node.name.value}"'
         decorators = [
             self.update_decorator(decorator=decorator, task_name=task_name)
-            if m.matches(decorator, celery_task_decorator)
+            if m.matches(
+                decorator, celery_app_task_decorator | celery_coalesced_task_decorator
+            )
             else decorator
             for decorator in updated_node.decorators
         ]

diff --git a/tests/test_command_componentize.py b/tests/test_command_componentize.py
@@ -93,6 +93,19 @@ def some_other_function(
                 topic: str, msg: bytes | dict, attributes: dict[str, str] | None = None
             ) -> None:
                 pass
+
+            @coalesced_task(
+                max_retries=10,
+                default_retry_delay=60,
+                queue=settings.TASK_QUEUE_LOW_LATENCY_TRACKING,
+            )
+            @app.some.other.decorator(
+                foo="bar",
+            )
+            def some_third_function(
+                topic: str, msg: bytes | dict, attributes: dict[str, str] | None = None
+            ) -> None:
+                pass
             """,
             """\
             @app.task(
@@ -122,6 +135,20 @@ def some_other_function(
                 topic: str, msg: bytes | dict, attributes: dict[str, str] | None = None
             ) -> None:
                 pass
+
+            @coalesced_task(
+                name="project.app.tasks.some_third_function",
+                max_retries=10,
+                default_retry_delay=60,
+                queue=settings.TASK_QUEUE_LOW_LATENCY_TRACKING,
+            )
+            @app.some.other.decorator(
+                foo="bar",
+            )
+            def some_third_function(
+                topic: str, msg: bytes | dict, attributes: dict[str, str] | None = None
+            ) -> None:
+                pass
             """,
         ),
     ],