Repository: scoutapp/scout_apm_python
Pull request #361 (2019-11-05 15:40:31)

Use wrapt for decorator in Django tests
It's shorter and simpler to read.


Patch:
diff --git a/src/scout_apm/instruments/elasticsearch.py b/src/scout_apm/instruments/elasticsearch.py
index b255134..4450e04 100644
--- a/src/scout_apm/instruments/elasticsearch.py
+++ b/src/scout_apm/instruments/elasticsearch.py
@@ -70,11 +70,11 @@ class Instrument(object):
             logger.info("Elasticsearch instruments are not installed. Skipping.")
             return False
 
-        self.__class__.installed = False
-
         self.uninstrument_client()
         self.uninstrument_transport()
 
+        self.__class__.installed = False
+
     def instrument_client(self):
         for name in self.__class__.CLIENT_METHODS:
             method = getattr(Elasticsearch, name, None)
diff --git a/src/scout_apm/instruments/pymongo.py b/src/scout_apm/instruments/pymongo.py
index 0119895..2e53c5d 100644
--- a/src/scout_apm/instruments/pymongo.py
+++ b/src/scout_apm/instruments/pymongo.py
@@ -3,10 +3,14 @@ from __future__ import absolute_import, division, print_function, unicode_litera
 
 import logging
 
-# Used in the exec() call below.
-import wrapt  # noqa: F401
+import wrapt
 
-from scout_apm.core.tracked_request import TrackedRequest  # noqa: F401
+from scout_apm.core.tracked_request import TrackedRequest
+
+try:
+    from pymongo.collection import Collection
+except ImportError:
+    Collection = None
 
 logger = logging.getLogger(__name__)
 
@@ -49,9 +53,7 @@ class Instrument(object):
     ]
 
     def installable(self):
-        try:
-            from pymongo.collection import Collection  # noqa: F401
-        except ImportError:
+        if Collection is None:
             logger.info("Unable to import for PyMongo instruments")
             return False
         if self.installed:
@@ -66,44 +68,10 @@ class Instrument(object):
 
         self.__class__.installed = True
 
-        try:
-            from pymongo.collection import Collection  # noqa: F401
-        # There is no way the import can fail if self.installable() succeeded.
-        except ImportError:  # pragma: no cover
-            logger.info(
-                "Unable to import for PyMongo instruments. Instrument install failed."
-            )
-            return False
-
-        for method_str in self.__class__.PYMONGO_METHODS:
-            try:
-                code_str = """
-@wrapt.decorator
-def wrapped_{method_str}(wrapped, instance, args, kwargs):
-    tracked_request = TrackedRequest.instance()
-    name = '/'.join(['MongoDB', instance.name, '{camel_name}'])
-    span = tracked_request.start_span(operation=name, ignore_children=True)
-    span.tag('name', instance.name)
-
-    try:
-        return wrapped(*args, **kwargs)
-    finally:
-        tracked_request.stop_span()
-
-Collection.{method_str} = wrapped_{method_str}(Collection.{method_str})
-""".format(
-                    method_str=method_str,
-                    camel_name="".join(c.title() for c in method_str.split("_")),
-                )
-
-                exec(code_str)
-                logger.info("Instrumented PyMongo Collection.%s", method_str)
-
-            except Exception as e:
-                logger.warning(
-                    "Unable to instrument for PyMongo Collection.%s: %r", method_str, e
-                )
-                return False
+        for name in self.__class__.PYMONGO_METHODS:
+            method = getattr(Collection, name, None)
+            if method is not None:
+                setattr(Collection, name, wrap_collection_method(method))
         return True
 
     def uninstall(self):
@@ -111,9 +79,23 @@ Collection.{method_str} = wrapped_{method_str}(Collection.{method_str})
             logger.info("PyMongo instruments are not installed. Skipping.")
             return False
 
+        for name in self.__class__.PYMONGO_METHODS:
+            method = getattr(Collection, name, None)
+            if method is not None:
+                setattr(Collection, name, method.__wrapped__)
+
         self.__class__.installed = False
 
-        from pymongo.collection import Collection
 
-        for method_str in self.__class__.PYMONGO_METHODS:
-            setattr(Collection, method_str, getattr(Collection, method_str).__wrapped__)
+@wrapt.decorator
+def wrap_collection_method(wrapped, instance, args, kwargs):
+    tracked_request = TrackedRequest.instance()
+    camel_name = "".join(c.title() for c in wrapped.__name__.split("_"))
+    operation = "MongoDB/{}.{}".format(instance.name, camel_name)
+    span = tracked_request.start_span(operation=operation, ignore_children=True)
+    span.tag("name", instance.name)
+
+    try:
+        return wrapped(*args, **kwargs)
+    finally:
+        tracked_request.stop_span()


Test patch:
diff --git a/tests/integration/django_app.py b/tests/integration/django_app.py
index bdc8e81..728fbf3 100644
--- a/tests/integration/django_app.py
+++ b/tests/integration/django_app.py
@@ -1,9 +1,8 @@
 # coding: utf-8
 from __future__ import absolute_import, division, print_function, unicode_literals
 
-from functools import wraps
-
 import django
+import wrapt
 from django.conf import settings
 from django.template.response import TemplateResponse
 
@@ -135,14 +134,11 @@ def template(request):
     return HttpResponse(template.render(context))
 
 
-def exclaimify_template_response_name(func):
-    @wraps(func)
-    def wrapper(*args, **kwargs):
-        response = func(*args, **kwargs)
-        response.context_data["name"] = response.context_data["name"] + "!!"
-        return response
-
-    return wrapper
+@wrapt.decorator
+def exclaimify_template_response_name(wrapped, instance, args, kwargs):
+    response = wrapped(*args, **kwargs)
+    response.context_data["name"] = response.context_data["name"] + "!!"
+    return response
 
 
 @exclaimify_template_response_name
diff --git a/tests/integration/instruments/test_elasticsearch.py b/tests/integration/instruments/test_elasticsearch.py
index b22bfac..c2b9616 100644
--- a/tests/integration/instruments/test_elasticsearch.py
+++ b/tests/integration/instruments/test_elasticsearch.py
@@ -170,7 +170,10 @@ def test_instrument_client_install_missing_attribute():
         "scout_apm.instruments.elasticsearch.Elasticsearch"
     ) as mock_elasticsearch:
         del mock_elasticsearch.bulk
-    instrument.instrument_client()  # no crash
+    try:
+        instrument.instrument_client()  # no crash
+    finally:
+        instrument.uninstrument_client()
 
 
 @mock.patch(
diff --git a/tests/integration/instruments/test_pymongo.py b/tests/integration/instruments/test_pymongo.py
index 8204b8c..207dc90 100644
--- a/tests/integration/instruments/test_pymongo.py
+++ b/tests/integration/instruments/test_pymongo.py
@@ -9,11 +9,9 @@ import pytest
 
 from scout_apm.instruments.pymongo import Instrument
 from tests.compat import mock
-from tests.tools import pretend_package_unavailable
 
 # e.g. export MONGODB_URL="mongodb://localhost:27017/"
 MONGODB_URL = os.environ.get("MONGODB_URL")
-
 skip_if_mongodb_not_running = pytest.mark.skipif(
     MONGODB_URL is None, reason="MongoDB isn't available"
 )
@@ -57,22 +55,22 @@ def test_installable():
 
 
 def test_installable_no_pymongo_module():
-    with pretend_package_unavailable("pymongo"):
+    with mock.patch("scout_apm.instruments.pymongo.Collection", new=None):
         assert not instrument.installable()
 
 
 def test_install_no_pymongo_module():
-    with pretend_package_unavailable("pymongo"):
+    with mock.patch("scout_apm.instruments.pymongo.Collection", new=None):
         assert not instrument.install()
 
 
-@mock.patch("scout_apm.instruments.pymongo.wrapt.decorator", side_effect=RuntimeError)
-def test_install_failure(mock_decorator):
+def test_install_missing_attribute():
+    with mock.patch("scout_apm.instruments.pymongo.Collection") as mock_collection:
+        del mock_collection.aggregate
     try:
-        assert not instrument.install()  # doesn't crash
+        instrument.install()  # no crash
     finally:
-        # Currently installed = True even if installing failed.
-        Instrument.installed = False
+        instrument.uninstall()
 
 
 def test_install_is_idempotent():