Add profile_trace testing

newrelic/api/profile_trace.py

@@ -13,31 +13,27 @@
 # limitations under the License.
 
 import functools
-import sys
 import os
+import sys
 
-from newrelic.packages import six
-
-from newrelic.api.time_trace import current_trace
+from newrelic import __file__ as AGENT_PACKAGE_FILE
 from newrelic.api.function_trace import FunctionTrace
-from newrelic.common.object_wrapper import FunctionWrapper, wrap_object
+from newrelic.api.time_trace import current_trace
 from newrelic.common.object_names import callable_name
+from newrelic.common.object_wrapper import FunctionWrapper, wrap_object
+from newrelic.packages import six
 
-from newrelic import __file__ as AGENT_PACKAGE_FILE
-AGENT_PACKAGE_DIRECTORY = os.path.dirname(AGENT_PACKAGE_FILE) + '/'
+AGENT_PACKAGE_DIRECTORY = os.path.dirname(AGENT_PACKAGE_FILE) + "/"
 
 
 class ProfileTrace(object):
-
     def __init__(self, depth):
         self.function_traces = []
         self.maximum_depth = depth
         self.current_depth = 0
 
-    def __call__(self, frame, event, arg):
-
-        if event not in ['call', 'c_call', 'return', 'c_return',
-                'exception', 'c_exception']:
+    def __call__(self, frame, event, arg):  # pragma: no cover
+        if event not in ["call", "c_call", "return", "c_return", "exception", "c_exception"]:
             return
 
         parent = current_trace()
@@ -49,8 +45,7 @@ def __call__(self, frame, event, arg):
         # coroutine systems based on greenlets so don't run
         # if we detect may be using greenlets.
 
-        if (hasattr(sys, '_current_frames') and
-                parent.thread_id not in sys._current_frames()):
+        if hasattr(sys, "_current_frames") and parent.thread_id not in sys._current_frames():
             return
 
         co = frame.f_code
@@ -84,7 +79,7 @@ def _callable():
                 except Exception:
                     pass
 
-        if event in ['call', 'c_call']:
+        if event in ["call", "c_call"]:
             # Skip the outermost as we catch that with the root
             # function traces for the profile trace.
 
@@ -100,19 +95,17 @@ def _callable():
                 self.function_traces.append(None)
                 return
 
-            if event == 'call':
+            if event == "call":
                 func = _callable()
                 if func:
                     name = callable_name(func)
                 else:
-                    name = '%s:%s#%s' % (func_filename, func_name,
-                            func_line_no)
+                    name = "%s:%s#%s" % (func_filename, func_name, func_line_no)
             else:
                 func = arg
                 name = callable_name(arg)
                 if not name:
-                    name = '%s:@%s#%s' % (func_filename, func_name,
-                            func_line_no)
+                    name = "%s:@%s#%s" % (func_filename, func_name, func_line_no)
 
             function_trace = FunctionTrace(name=name, parent=parent)
             function_trace.__enter__()
@@ -127,7 +120,7 @@ def _callable():
             self.function_traces.append(function_trace)
             self.current_depth += 1
 
-        elif event in ['return', 'c_return', 'c_exception']:
+        elif event in ["return", "c_return", "c_exception"]:
             if not self.function_traces:
                 return
 
@@ -143,9 +136,7 @@ def _callable():
                     self.current_depth -= 1
 
 
-def ProfileTraceWrapper(wrapped, name=None, group=None, label=None,
-        params=None, depth=3):
-
+def ProfileTraceWrapper(wrapped, name=None, group=None, label=None, params=None, depth=3):
     def wrapper(wrapped, instance, args, kwargs):
         parent = current_trace()
 
@@ -192,7 +183,7 @@ def wrapper(wrapped, instance, args, kwargs):
             _params = params
 
         with FunctionTrace(_name, _group, _label, _params, parent=parent, source=wrapped):
-            if not hasattr(sys, 'getprofile'):
+            if not hasattr(sys, "getprofile"):
                 return wrapped(*args, **kwargs)
 
             profiler = sys.getprofile()
@@ -212,11 +203,8 @@ def wrapper(wrapped, instance, args, kwargs):
 
 
 def profile_trace(name=None, group=None, label=None, params=None, depth=3):
-    return functools.partial(ProfileTraceWrapper, name=name,
-            group=group, label=label, params=params, depth=depth)
+    return functools.partial(ProfileTraceWrapper, name=name, group=group, label=label, params=params, depth=depth)
 
 
-def wrap_profile_trace(module, object_path, name=None,
-        group=None, label=None, params=None, depth=3):
-    return wrap_object(module, object_path, ProfileTraceWrapper,
-            (name, group, label, params, depth))
+def wrap_profile_trace(module, object_path, name=None, group=None, label=None, params=None, depth=3):
+    return wrap_object(module, object_path, ProfileTraceWrapper, (name, group, label, params, depth))

tests/agent_features/test_profile_trace.py

@@ -0,0 +1,88 @@
+# Copyright 2010 New Relic, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from testing_support.validators.validate_transaction_metrics import (
+    validate_transaction_metrics,
+)
+
+from newrelic.api.background_task import background_task
+from newrelic.api.profile_trace import ProfileTraceWrapper, profile_trace
+
+
+def test_profile_trace_wrapper():
+    def _test():
+        def nested_fn():
+            pass
+
+        nested_fn()
+
+    wrapped_test = ProfileTraceWrapper(_test)
+    wrapped_test()
+
+
+@validate_transaction_metrics("test_profile_trace:test_profile_trace_empty_args", background_task=True)
+@background_task()
+def test_profile_trace_empty_args():
+    @profile_trace()
+    def _test():
+        pass
+
+    _test()
+
+
+_test_profile_trace_defined_args_scoped_metrics = [("Custom/TestTrace", 1)]
+
+
+@validate_transaction_metrics(
+    "test_profile_trace:test_profile_trace_defined_args",
+    scoped_metrics=_test_profile_trace_defined_args_scoped_metrics,
+    background_task=True,
+)
+@background_task()
+def test_profile_trace_defined_args():
+    @profile_trace(name="TestTrace", group="Custom", label="Label", params={"key": "value"}, depth=7)
+    def _test():
+        pass
+
+    _test()
+
+
+_test_profile_trace_callable_args_scoped_metrics = [("Function/TestProfileTrace", 1)]
+
+
+@validate_transaction_metrics(
+    "test_profile_trace:test_profile_trace_callable_args",
+    scoped_metrics=_test_profile_trace_callable_args_scoped_metrics,
+    background_task=True,
+)
+@background_task()
+def test_profile_trace_callable_args():
+    def name_callable():
+        return "TestProfileTrace"
+
+    def group_callable():
+        return "Function"
+
+    def label_callable():
+        return "HSM"
+
+    def params_callable():
+        return {"account_id": "12345"}
+
+    @profile_trace(name=name_callable, group=group_callable, label=label_callable, params=params_callable, depth=0)
+    def _test():
+        pass
+
+    _test()