Fix Starlette instrumentation after 0.20.1 release (#560)

* Fix Starlette instrumentation after 0.20.1 release Co-authored-by: Timothy Pansino <[email protected]> Co-authored-by: Uma Annamalai <[email protected]> * Fix the hasattr gate for Starlette instrumentation Co-authored-by: Timothy Pansino <[email protected]> * Fix basehttp test for Python 3.7/Starlette 0.20.1 * Clean up tests in Starlette * Remove import of "version_info" Co-authored-by: Timothy Pansino <[email protected]> * Update tests/framework_starlette/test_bg_tasks.py Co-authored-by: Timothy Pansino <[email protected]> Co-authored-by: Uma Annamalai <[email protected]> Co-authored-by: Timothy Pansino <[email protected]>
newrelic · Jun 2, 2022 · 6c8e8a9 · 6c8e8a9
1 parent 1ea34ed
commit 6c8e8a9
Show file tree

Hide file tree

Showing 6 changed files with 89 additions and 30 deletions.
diff --git a/newrelic/config.py b/newrelic/config.py
@@ -2467,6 +2467,11 @@ def _process_module_builtin_defaults():
         "newrelic.hooks.framework_starlette",
         "instrument_starlette_middleware_errors",
     )
+    _process_module_definition(
+        "starlette.middleware.exceptions",
+        "newrelic.hooks.framework_starlette",
+        "instrument_starlette_middleware_exceptions",
+    )
     _process_module_definition(
         "starlette.exceptions",
         "newrelic.hooks.framework_starlette",

diff --git a/newrelic/hooks/framework_starlette.py b/newrelic/hooks/framework_starlette.py
@@ -241,14 +241,26 @@ def instrument_starlette_middleware_errors(module):
     wrap_function_wrapper(module, "ServerErrorMiddleware.debug_response", wrap_exception_handler)
 
 
-def instrument_starlette_exceptions(module):
+def instrument_starlette_middleware_exceptions(module):
     wrap_function_wrapper(module, "ExceptionMiddleware.__call__", error_middleware_wrapper)
 
     wrap_function_wrapper(module, "ExceptionMiddleware.http_exception", wrap_exception_handler)
 
     wrap_function_wrapper(module, "ExceptionMiddleware.add_exception_handler", wrap_add_exception_handler)
 
 
+def instrument_starlette_exceptions(module):
+    # ExceptionMiddleware was moved to starlette.middleware.exceptions, need to check
+    # that it isn't being imported through a deprecation and double wrapped.
+    if not hasattr(module, "__deprecated__"):
+
+        wrap_function_wrapper(module, "ExceptionMiddleware.__call__", error_middleware_wrapper)
+
+        wrap_function_wrapper(module, "ExceptionMiddleware.http_exception", wrap_exception_handler)
+
+        wrap_function_wrapper(module, "ExceptionMiddleware.add_exception_handler", wrap_add_exception_handler)
+
+
 def instrument_starlette_background_task(module):
     wrap_function_wrapper(module, "BackgroundTask.__call__", wrap_background_method)
 

diff --git a/tests/framework_starlette/_test_bg_tasks.py b/tests/framework_starlette/_test_bg_tasks.py
@@ -52,7 +52,7 @@ async def async_bg_task():
     pass
 
 
-async def sync_bg_task():
+def sync_bg_task():
     pass
 
 

diff --git a/tests/framework_starlette/test_application.py b/tests/framework_starlette/test_application.py
@@ -25,6 +25,7 @@
 from newrelic.common.object_names import callable_name
 from testing_support.validators.validate_code_level_metrics import validate_code_level_metrics
 
+starlette_version = tuple(int(x) for x in starlette.__version__.split("."))
 
 @pytest.fixture(scope="session")
 def target_application():
@@ -34,10 +35,18 @@ def target_application():
 
 
 FRAMEWORK_METRIC = ("Python/Framework/Starlette/%s" % starlette.__version__, 1)
-DEFAULT_MIDDLEWARE_METRICS = [
-    ("Function/starlette.middleware.errors:ServerErrorMiddleware.__call__", 1),
-    ("Function/starlette.exceptions:ExceptionMiddleware.__call__", 1),
-]
+
+if starlette_version >= (0, 20, 1):
+    DEFAULT_MIDDLEWARE_METRICS = [
+        ("Function/starlette.middleware.errors:ServerErrorMiddleware.__call__", 1),
+        ("Function/starlette.middleware.exceptions:ExceptionMiddleware.__call__", 1),
+    ]
+else:
+    DEFAULT_MIDDLEWARE_METRICS = [
+        ("Function/starlette.middleware.errors:ServerErrorMiddleware.__call__", 1),
+        ("Function/starlette.exceptions:ExceptionMiddleware.__call__", 1),
+    ]
+
 MIDDLEWARE_METRICS = [
     ("Function/_test_application:middleware_factory.<locals>.middleware", 2),
     ("Function/_test_application:middleware_decorator", 1),
@@ -71,17 +80,27 @@ def test_application_non_async(target_application, app_name):
     response = app.get("/non_async")
     assert response.status == 200
 
+# Starting in Starlette v0.20.1, the ExceptionMiddleware class
+# has been moved to the starlette.middleware.exceptions from
+# starlette.exceptions
+version_tweak_string = ".middleware" if starlette_version >= (0, 20, 1) else ""
 
-@pytest.mark.parametrize(
-    "app_name, transaction_name",
+DEFAULT_MIDDLEWARE_METRICS = [
+    ("Function/starlette.middleware.errors:ServerErrorMiddleware.__call__", 1),
+    ("Function/starlette%s.exceptions:ExceptionMiddleware.__call__" % version_tweak_string, 1),
+]
+
+middleware_test = (
+    ("no_error_handler", "starlette%s.exceptions:ExceptionMiddleware.__call__" % version_tweak_string),
     (
-        ("no_error_handler", "starlette.exceptions:ExceptionMiddleware.__call__"),
-        (
-            "non_async_error_handler_no_middleware",
-            "starlette.exceptions:ExceptionMiddleware.__call__",
-        ),
+        "non_async_error_handler_no_middleware",
+        "starlette%s.exceptions:ExceptionMiddleware.__call__" % version_tweak_string,
     ),
 )
+
+@pytest.mark.parametrize(
+    "app_name, transaction_name", middleware_test,
+)
 def test_application_nonexistent_route(target_application, app_name, transaction_name):
     @validate_transaction_metrics(
         transaction_name,
@@ -244,19 +263,20 @@ def _test():
     _test()
 
 
-@pytest.mark.parametrize(
-    "app_name,scoped_metrics",
+middleware_test_exception = (
     (
-        (
-            "no_middleware",
-            [("Function/starlette.exceptions:ExceptionMiddleware.http_exception", 1)],
-        ),
-        (
-            "teapot_exception_handler_no_middleware",
-            [("Function/_test_application:teapot_handler", 1)],
-        ),
+        "no_middleware",
+        [("Function/starlette%s.exceptions:ExceptionMiddleware.http_exception" % version_tweak_string, 1)],
+    ),
+    (
+        "teapot_exception_handler_no_middleware",
+        [("Function/_test_application:teapot_handler", 1)],
     ),
 )
+
+@pytest.mark.parametrize(
+    "app_name,scoped_metrics", middleware_test_exception
+)
 def test_starlette_http_exception(target_application, app_name, scoped_metrics):
     @validate_transaction_errors(errors=["starlette.exceptions:HTTPException"])
     @validate_transaction_metrics(

diff --git a/tests/framework_starlette/test_bg_tasks.py b/tests/framework_starlette/test_bg_tasks.py
@@ -13,13 +13,17 @@
 # limitations under the License.
 
 import pytest
+import sys
 from testing_support.fixtures import validate_transaction_metrics
 from testing_support.validators.validate_transaction_count import (
     validate_transaction_count,
 )
 
+from starlette import __version__
+starlette_version = tuple(int(x) for x in __version__.split("."))
+
 try:
-    from starlette.middleware import Middleware
+    from starlette.middleware import Middleware     # Ignore Flake8 Error
 
     no_middleware = False
 except ImportError:
@@ -79,18 +83,35 @@ def _test():
 @skip_if_no_middleware
 @pytest.mark.parametrize("route", ["async", "sync"])
 def test_basehttp_style_middleware(target_application, route):
-    metrics = [
+    route_metrics = [("Function/_test_bg_tasks:run_%s_bg_task" % route, 1)]
+    old_metrics = [
         ("Function/_test_bg_tasks:%s_bg_task" % route, 1),
         ("Function/_test_bg_tasks:run_%s_bg_task" % route, 1),
     ]
 
-    @validate_transaction_metrics(
-        "_test_bg_tasks:run_%s_bg_task" % route, scoped_metrics=metrics
-    )
-    @validate_transaction_count(1)
     def _test():
         app = target_application["basehttp"]
         response = app.get("/" + route)
         assert response.status == 200
 
+    if starlette_version >= (0, 20, 1):
+        if sys.version_info[:2] > (3, 7):
+            _test = validate_transaction_metrics(
+                "_test_bg_tasks:run_%s_bg_task" % route, index=-2, scoped_metrics=route_metrics
+            )(_test)
+            _test = validate_transaction_metrics(
+                "_test_bg_tasks:%s_bg_task" % route, background_task=True
+            )(_test)
+            _test = validate_transaction_count(2)(_test)
+        else:   # Python <= 3.7 requires this specific configuration with starlette 0.20.1
+            _test = validate_transaction_metrics(
+                "_test_bg_tasks:run_%s_bg_task" % route, scoped_metrics=route_metrics
+            )(_test)
+            _test = validate_transaction_count(1)(_test)
+    else:
+        _test = validate_transaction_metrics(
+            "_test_bg_tasks:run_%s_bg_task" % route, scoped_metrics=old_metrics
+        )(_test)
+        _test = validate_transaction_count(1)(_test)
+
     _test()
diff --git a/tox.ini b/tox.ini
@@ -136,7 +136,7 @@ envlist =
     python-framework_pyramid-{py37,py38,py39,py310}-Pyramidmaster,
     python-framework_sanic-{py38,pypy3}-sanic{190301,1906,1812,1912,200904,210300},
     python-framework_sanic-{py36,py37,py38,py310,pypy3}-saniclatest,
-    python-framework_starlette-{py36,py310,pypy3}-starlette{0014,0015},
+    python-framework_starlette-{py36,py310,pypy3}-starlette{0014,0015,0019},
     python-framework_starlette-{py36,py37,py38,py39,py310,pypy3}-starlette{latest},
     python-framework_strawberry-{py37,py38,py39,py310}-strawberrylatest,
     libcurl-framework_tornado-{py36,py37,py38,py39,py310,pypy3}-tornado0600,
@@ -316,6 +316,7 @@ deps =
     framework_starlette: graphene<3
     framework_starlette-starlette0014: starlette<0.15
     framework_starlette-starlette0015: starlette<0.16
+    framework_starlette-starlette0019: starlette<0.20
     framework_starlette-starlettelatest: starlette
     ; Strawberry 0.95.0 is incompatible with Starlette 0.18.0, downgrade until future release
     framework_strawberry: starlette<0.18.0
-Original file line number
+Diff line change
@@ Expand Up / @@ -52,7 +52,7 @@ async def async_bg_task(): @@
         pass
-    async def sync_bg_task():
+    def sync_bg_task():
         pass
@@ Expand Down @@