Repository: scoutapp/scout_apm_python (python)

Issue:
Support Jinja 3.x
Jinja 3.x was released last week and it has broken some functionality within the agent. The following tests are currently failing:

```
tests/integration/instruments/test_jinja2_py36plus.py::test_async_render
tests/integration/instruments/test_jinja2_py36plus.py::test_async_render_name
tests/integration/instruments/test_urllib3.py::test_request - pytest.PytestUnraisableExceptionWarning
tests/integration/instruments/test_urllib3.py::test_request_no_absolute_url
```


Pull request description:
Instrument Jinja v3 render_async.
Fixes #646

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index b249935..7c36fcb 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,8 @@
 
 ### Fixed
 - Switch to storing TrackedRequest in a ContextVar for Python3.6+ rather than in a thread local.
+- Support instrumenting Jinja v3 ``render_async`` function.
+  ([Issue #646](https://github.com/scoutapp/scout_apm_python/issues/646))
 
 ## [2.21.1] 2021-07-27
 
diff --git a/src/scout_apm/instruments/jinja2.py b/src/scout_apm/instruments/jinja2.py
index 31970d5..042f2d6 100644
--- a/src/scout_apm/instruments/jinja2.py
+++ b/src/scout_apm/instruments/jinja2.py
@@ -34,7 +34,6 @@ have_patched_template_render_async = False
 
 
 def ensure_installed():
-    global have_patched_environment_init
     global have_patched_template_render
 
     logger.debug("Instrumenting Jinja2.")
@@ -43,9 +42,31 @@ def ensure_installed():
         logger.debug("Couldn't import jinja2.Template - probably not installed.")
         return
 
-    if not have_patched_environment_init:
+    instrument_render_async()
+
+    if not have_patched_template_render:
+        try:
+            Template.render = wrapped_render(Template.render)
+        except Exception as exc:
+            logger.warning(
+                "Failed to instrument jinja2.Template.render: %r", exc, exc_info=exc
+            )
+        else:
+            have_patched_template_render = True
+
+
+def instrument_render_async():
+    global have_patched_environment_init
+    global have_patched_template_render_async
+
+    if wrapped_render_async is None:
+        return
+
+    if not have_patched_environment_init and not hasattr(Template, "render_async"):
         try:
-            Environment.__init__ = wrapped_environment_init(Environment.__init__)
+            Environment.__init__ = wrapped_environment_init_jinja_v2(
+                Environment.__init__
+            )
         except Exception as exc:
             logger.warning(
                 "Failed to instrument jinja2.Environment.__init__: %r",
@@ -54,16 +75,17 @@ def ensure_installed():
             )
         else:
             have_patched_environment_init = True
-
-    if not have_patched_template_render:
+    elif hasattr(Template, "render_async") and not have_patched_template_render_async:
         try:
-            Template.render = wrapped_render(Template.render)
+            Template.render_async = wrapped_render_async(Template.render_async)
         except Exception as exc:
             logger.warning(
-                "Failed to instrument jinja2.Template.render: %r", exc, exc_info=exc
+                "Failed to instrument jinja2.Template.render_async: %r",
+                exc,
+                exc_info=exc,
             )
         else:
-            have_patched_template_render = True
+            have_patched_template_render_async = True
 
 
 @wrapt.decorator
@@ -75,12 +97,14 @@ def wrapped_render(wrapped, instance, args, kwargs):
 
 
 @wrapt.decorator
-def wrapped_environment_init(wrapped, instance, args, kwargs):
+def wrapped_environment_init_jinja_v2(wrapped, instance, args, kwargs):
     """
     Delayed wrapping of render_async(), since Template won't have this method
     until after jinja2.asyncsupport is imported, which since Jinja2 2.11.0 is
     done conditionally in Environment.__init__:
     https://github.com/pallets/jinja/issues/765
+
+    This is no longer needed since Jinja2 v3.0.0
     """
     global have_patched_template_render_async
     result = wrapped(*args, **kwargs)
diff --git a/tox.ini b/tox.ini
index c3f5225..ae6a53e 100644
--- a/tox.ini
+++ b/tox.ini
@@ -51,7 +51,7 @@ deps =
     hug>=2.5.1 ; python_version >= "3.5"
     httpretty<1 ; python_version < "3.5"
     httpretty!=1.1.1 ; python_version >= "3.5"
-    jinja2<3.0.0
+    jinja2
     ; nameko tests temporarily disabled
     ; Due to bad pinning of kombu only a past nameko version can be installed
     ; alongside the latest celery.


Test patch:
diff --git a/tests/integration/instruments/test_jinja2_py36plus.py b/tests/integration/instruments/test_jinja2_py36plus.py
index 478e3bb..da411f0 100644
--- a/tests/integration/instruments/test_jinja2_py36plus.py
+++ b/tests/integration/instruments/test_jinja2_py36plus.py
@@ -11,7 +11,7 @@ from tests.compat import mock
 from tests.integration.instruments.test_jinja2 import mock_not_attempted
 
 
-def test_ensure_installed_fail_no_render_async_attribute(caplog):
+def test_ensure_installed_render_async_v2(caplog):
     mock_template = mock.patch("scout_apm.instruments.jinja2.Template")
     with mock_not_attempted(), mock_template as mocked_template:
         del mocked_template.render_async
@@ -29,6 +29,29 @@ def test_ensure_installed_fail_no_render_async_attribute(caplog):
     ]
 
 
+def test_ensure_installed_fail_wrap_render_async(caplog):
+    mock_wrap_render_async = mock.patch(
+        "scout_apm.instruments.jinja2.wrapped_render_async"
+    )
+    with mock_not_attempted(), mock_wrap_render_async as mocked_wrap_render_async:
+        mocked_wrap_render_async.side_effect = ValueError("Boom!")
+        ensure_installed()
+
+    assert len(caplog.record_tuples) == 2
+    assert caplog.record_tuples[0] == (
+        "scout_apm.instruments.jinja2",
+        logging.DEBUG,
+        "Instrumenting Jinja2.",
+    )
+
+    logger, level, message = caplog.record_tuples[1]
+    assert logger == "scout_apm.instruments.jinja2"
+    assert level == logging.WARNING
+    assert message.startswith(
+        "Failed to instrument jinja2.Template.render_async: ValueError"
+    )
+
+
 @pytest.mark.asyncio
 async def test_async_render(tracked_request):
     ensure_installed()