Added tests for CursorWrapper patching

django_statsd/patches/db.py

@@ -7,6 +7,7 @@ from django_statsd.clients import statsd
 def key(db, attr):
     return 'db.%s.%s.%s' % (db.client.executable_name, db.alias, attr)
 
+
 def pre_django_1_6_cursorwrapper_getattr(self, attr):
     """
     The CursorWrapper is a pretty small wrapper around the cursor.
@@ -22,33 +23,39 @@ def pre_django_1_6_cursorwrapper_getattr(self, attr):
             return wrap(getattr(self.cursor, attr), key(self.db, attr))
         return getattr(self.cursor, attr)
 
+
+def patched_execute(orig_execute, self, *args, **kwargs):
+    with statsd.timer(key(self.db, 'execute')):
+        return orig_execute(self, *args, **kwargs)
+
+
+def patched_executemany(orig_executemany, self, *args, **kwargs):
+    with statsd.timer(key(self.db, 'executemany')):
+        return orig_executemany(self, *args, **kwargs)
+
+
+def patched_callproc(orig_callproc, self, *args, **kwargs):
+    with statsd.timer(key(self.db, 'callproc')):
+        return orig_callproc(self, *args, **kwargs)
+
+
 def patch():
     """
-    The CursorWrapper is a pretty small wrapper around the cursor.
-    If you are NOT in debug mode, this is the wrapper that's used.
-    Sadly if it's in debug mode, we get a different wrapper for version earlier than 1.6.
+    The CursorWrapper is a pretty small wrapper around the cursor.  If
+    you are NOT in debug mode, this is the wrapper that's used.  Sadly
+    if it's in debug mode, we get a different wrapper for version
+    earlier than 1.6.
     """
 
-    def execute(orig_execute, self, *args, **kwargs):
-        with statsd.timer(key(self.db, 'execute')):
-            return orig_execute(self, *args, **kwargs)
-
-    def executemany(orig_executemany, self, *args, **kwargs):
-        with statsd.timer(key(self.db, 'executemany')):
-            return orig_executemany(self, *args, **kwargs)
-
-    def callproc(orig_callproc, self, *args, **kwargs):
-        with statsd.timer(key(self.db, 'callproc')):
-            return orig_callproc(self, *args, **kwargs)
-
     if django.VERSION > (1, 6):
         # In 1.6+ util.CursorDebugWrapper just makes calls to CursorWrapper
         # As such, we only need to instrument CursorWrapper.
         # Instrumenting both will result in duplicated metrics
-        patch_method(util.CursorWrapper, 'execute')(execute)
-        patch_method(util.CursorWrapper, 'executemany')(executemany)
-        patch_method(util.CursorWrapper, 'callproc')(callproc)
+        patch_method(util.CursorWrapper, 'execute')(patched_execute)
+        patch_method(util.CursorWrapper, 'executemany')(patched_executemany)
+        patch_method(util.CursorWrapper, 'callproc')(patched_callproc)
     else:
         util.CursorWrapper.__getattr__ = pre_django_1_6_cursorwrapper_getattr
-        patch_method(util.CursorDebugWrapper, 'execute')(execute)
-        patch_method(util.CursorDebugWrapper, 'executemany')(executemany)
+        patch_method(util.CursorDebugWrapper, 'execute')(patched_execute)
+        patch_method(
+            util.CursorDebugWrapper, 'executemany')(patched_executemany)

django_statsd/tests.py

@@ -5,7 +5,9 @@ import sys
 from django.conf import settings
 from nose.exc import SkipTest
 from nose import tools as nose_tools
+from unittest import skipUnless
 
+from django import VERSION
 from django.core.urlresolvers import reverse
 from django.http import HttpResponse, HttpResponseForbidden
 from django.test import TestCase
@@ -15,7 +17,7 @@ from django.utils import unittest
 
 import mock
 from nose.tools import eq_
-from django_statsd.clients import get_client
+from django_statsd.clients import get_client, statsd
 from django_statsd.patches import utils
 from django_statsd import middleware
 
@@ -464,3 +466,90 @@ class TestPatchMethod(TestCase):
         self.assertEqual(obj.badfn(1, d=2), ((1,), {'d': 2}))
         self.assertEqual(obj.badfn(1, d=2), ((1,), {'d': 2}))
         self.assertEqual(obj.badfn(1, 2, c=1, d=2), ((1, 2), {'c': 1, 'd': 2}))
+
+
+class TestCursorWrapperPatching(TestCase):
+
+    def test_patched_callproc_calls_timer(self):
+        from django_statsd.patches.db import patched_callproc
+        with mock.patch.object(statsd, 'timer') as timer:
+            db = mock.Mock(executable_name='name', alias='alias')
+            instance = mock.Mock(db=db)
+            patched_callproc(lambda *args, **kwargs: None, instance)
+            self.assertEqual(timer.call_count, 1)
+
+    def test_patched_execute_calls_timer(self):
+        from django_statsd.patches.db import patched_execute
+        with mock.patch.object(statsd, 'timer') as timer:
+            db = mock.Mock(executable_name='name', alias='alias')
+            instance = mock.Mock(db=db)
+            patched_execute(lambda *args, **kwargs: None, instance)
+            self.assertEqual(timer.call_count, 1)
+
+    def test_patched_executemany_calls_timer(self):
+        from django_statsd.patches.db import patched_executemany
+        with mock.patch.object(statsd, 'timer') as timer:
+            db = mock.Mock(executable_name='name', alias='alias')
+            instance = mock.Mock(db=db)
+            patched_executemany(lambda *args, **kwargs: None, instance)
+            self.assertEqual(timer.call_count, 1)
+
+    @mock.patch(
+        'django_statsd.patches.db.pre_django_1_6_cursorwrapper_getattr')
+    @mock.patch('django_statsd.patches.db.patched_executemany')
+    @mock.patch('django_statsd.patches.db.patched_execute')
+    @mock.patch('django.db.backends.util.CursorDebugWrapper')
+    @skipUnless(VERSION < (1, 6, 0), "CursorWrapper Patching for Django<1.6")
+    def test_cursorwrapper_patching(self,
+                                    CursorDebugWrapper,
+                                    execute,
+                                    executemany,
+                                    _getattr):
+        try:
+            from django.db.backends import util
+
+            # We need to patch CursorWrapper like this because setting
+            # __getattr__ on Mock instances raises AttributeError.
+            class CursorWrapper(object):
+                pass
+
+            _CursorWrapper = util.CursorWrapper
+            util.CursorWrapper = CursorWrapper
+
+            from django_statsd.patches.db import patch
+            execute.__name__ = 'execute'
+            executemany.__name__ = 'executemany'
+            _getattr.__name__ = '_getattr'
+            execute.return_value = 'execute'
+            executemany.return_value = 'executemany'
+            _getattr.return_value = 'getattr'
+            patch()
+
+            self.assertEqual(CursorDebugWrapper.execute(), 'execute')
+            self.assertEqual(CursorDebugWrapper.executemany(), 'executemany')
+            self.assertEqual(CursorWrapper.__getattr__(), 'getattr')
+        finally:
+            util.CursorWrapper = _CursorWrapper
+
+    @mock.patch('django_statsd.patches.db.patched_callproc')
+    @mock.patch('django_statsd.patches.db.patched_executemany')
+    @mock.patch('django_statsd.patches.db.patched_execute')
+    @mock.patch('django.db.backends.util.CursorWrapper')
+    @skipUnless(VERSION >= (1, 6, 0), "CursorWrapper Patching for Django>=1.6")
+    def test_cursorwrapper_patching16(self,
+                                      CursorWrapper,
+                                      execute,
+                                      executemany,
+                                      callproc):
+        from django_statsd.patches.db import patch
+        execute.__name__ = 'execute'
+        executemany.__name__ = 'executemany'
+        callproc.__name__ = 'callproc'
+        execute.return_value = 'execute'
+        executemany.return_value = 'executemany'
+        callproc.return_value = 'callproc'
+        patch()
+
+        self.assertEqual(CursorWrapper.execute(), 'execute')
+        self.assertEqual(CursorWrapper.executemany(), 'executemany')
+        self.assertEqual(CursorWrapper.callproc(), 'callproc')