Fix remaining tests and add custom task tests

Author: lrafeei

newrelic/hooks/application_celery.py

@@ -264,51 +264,6 @@ def run(self, *args, **kwargs):
     wrapped_task.__module__ = CeleryTaskWrapper.__module__
 
     return wrapped_task
-            
-
-# # This will not work with the current version of Celery
-# # This only gets called during the async execution of a task
-# # and the task is wrapped later in the process to accomodate
-# # custom task classes.
-# def wrap_Celery_send_task(wrapped, instance, args, kwargs):
-#     transaction = current_transaction()
-#     if not transaction:
-#         return wrapped(*args, **kwargs)
-
-#     # Merge distributed tracing headers into outgoing task headers
-#     try:
-#         dt_headers = MessageTrace.generate_request_headers(transaction)
-#         original_headers = kwargs.get("headers", None)
-#         if dt_headers:
-#             if not original_headers:
-#                 kwargs["headers"] = dict(dt_headers)
-#             else:
-#                 kwargs["headers"] = dt_headers = dict(dt_headers)
-#                 dt_headers.update(dict(original_headers))
-#     except Exception:
-#         pass
-
-#     return wrapped(*args, **kwargs)
-
-
-def wrap_worker_optimizations(wrapped, instance, args, kwargs):
-    # Attempt to uninstrument BaseTask before stack protection is installed or uninstalled
-    try:
-        from celery.app.task import BaseTask
-
-        if isinstance(BaseTask.__call__, _NRBoundFunctionWrapper):
-            BaseTask.__call__ = BaseTask.__call__.__wrapped__
-    except Exception:
-        BaseTask = None
-
-    # Allow metaprogramming to run
-    result = wrapped(*args, **kwargs)
-
-    # Rewrap finalized BaseTask
-    if BaseTask:  # Ensure imports succeeded
-        BaseTask.__call__ = CeleryTaskWrapper(BaseTask.__call__)
-
-    return result
 
 
 def instrument_celery_local(module):
@@ -319,11 +274,6 @@ def instrument_celery_local(module):
         module.Proxy.__call__ = CeleryTaskWrapper(module.Proxy.__call__)
 
 
-# def instrument_celery_app_base(module):
-#     if hasattr(module, "Celery") and hasattr(module.Celery, "send_task"):
-#         wrap_function_wrapper(module, "Celery.send_task", wrap_Celery_send_task)
-
-
 def instrument_celery_worker(module):
     if hasattr(module, "process_initializer"):
         # We try and force activation of the agent before 
@@ -367,13 +317,6 @@ def force_agent_shutdown(*args, **kwargs):
 
 
 def instrument_celery_app_trace(module):
-    # Uses same wrapper for setup and reset worker optimizations to prevent patching and unpatching from removing wrappers
-    if hasattr(module, "setup_worker_optimizations"):
-        wrap_function_wrapper(module, "setup_worker_optimizations", wrap_worker_optimizations)
-
-    if hasattr(module, "reset_worker_optimizations"):
-        wrap_function_wrapper(module, "reset_worker_optimizations", wrap_worker_optimizations)
-
     if hasattr(module, "build_tracer"):
         wrap_function_wrapper(module, "build_tracer", wrap_build_tracer)
 

tests/application_celery/_target_application.py

@@ -12,8 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from celery import Celery, shared_task
-from testing_support.validators.validate_distributed_trace_accepted import validate_distributed_trace_accepted
+from celery import Celery, Task, shared_task
 
 from newrelic.api.transaction import current_transaction
 
@@ -26,12 +25,36 @@
     broker_heartbeat=0,
 )
 
+class CustomCeleryTaskWithSuper(Task):
+    def __call__(self, *args, **kwargs):
+        transaction = current_transaction()
+        if transaction:
+            transaction.add_custom_attribute("custom_task_attribute", "Called with super")
+        return super().__call__(*args, **kwargs)
+
+class CustomCeleryTaskWithRun(Task):
+    def __call__(self, *args, **kwargs):
+        transaction = current_transaction()
+        if transaction:
+            transaction.add_custom_attribute("custom_task_attribute", "Called with run")
+        return self.run(*args, **kwargs)
+
 
 @app.task
 def add(x, y):
     return x + y
 
 
+@app.task(base=CustomCeleryTaskWithSuper)
+def add_with_super(x, y):
+    return x + y
+
+
+@app.task(base=CustomCeleryTaskWithRun)
+def add_with_run(x, y):
+    return x + y
+
+
 @app.task
 def tsum(nums):
     return sum(nums)
@@ -45,13 +68,3 @@ def nested_add(x, y):
 @shared_task
 def shared_task_add(x, y):
     return x + y
-
-
-@app.task
-@validate_distributed_trace_accepted(transport_type="AMQP")
-def assert_dt():
-    # Basic checks for DT delegated to task
-    txn = current_transaction()
-    assert txn, "No transaction active."
-    assert txn.name == "_target_application.assert_dt", f"Transaction name does not match: {txn.name}"
-    return 1

tests/application_celery/test_application.py

@@ -24,11 +24,9 @@
 @validate_transaction_metrics(
     name="test_application:test_celery_task_as_function_trace",
     scoped_metrics=[("Function/_target_application.add", 1)],
-    # scoped_metrics=[("Function/_target_application.add", 2)],
     background_task=True,
 )
 @validate_code_level_metrics("_target_application", "add")
-# @validate_code_level_metrics("_target_application", "add", count=2)
 @background_task()
 def test_celery_task_as_function_trace():
     """