Add utilities to inject and monitor functions. Change DB tracking to use new injection method on BaseDatabaseWrapper.cursor

1 files changed, 136 insertions, 1 deletions

diff --git a/debug_toolbar/tests/tests.py b/debug_toolbar/tests/tests.py
index cdc1947..476802c 100644
--- a/debug_toolbar/tests/tests.py
+++ b/debug_toolbar/tests/tests.py
@@ -1,6 +1,7 @@
 from debug_toolbar.middleware import DebugToolbarMiddleware
 from debug_toolbar.panels.sql import SQLDebugPanel
 from debug_toolbar.toolbar.loader import DebugToolbar
+from debug_toolbar.utils.tracking import pre_dispatch, post_dispatch, callbacks
 
 from django.contrib.auth.models import User
 from django.test import TestCase
@@ -36,4 +37,138 @@ class SQLPanelTestCase(BaseTestCase):
         self.assertTrue('sql' in query[1])
         self.assertTrue('duration' in query[1])
         self.assertTrue('stacktrace' in query[1])
-        
\ No newline at end of file
+
+def module_func(*args, **kwargs):
+    """Used by dispatch tests"""
+    return 'blah'
+
+class TrackingTestCase(BaseTestCase):
+    @classmethod
+    def class_method(cls, *args, **kwargs):
+        return 'blah'
+
+    def class_func(self, *args, **kwargs):
+        """Used by dispatch tests"""
+        return 'blah'
+    
+    def test_pre_hook(self):
+        foo = {}
+        
+        @pre_dispatch(module_func)
+        def test(**kwargs):
+            foo.update(kwargs)
+            
+        self.assertTrue(hasattr(module_func, '__wrapped__'))
+        self.assertEquals(len(callbacks['before']), 1)
+        
+        module_func('hi', foo='bar')
+        
+        self.assertTrue('sender' in foo, foo)
+        # best we can do
+        self.assertEquals(foo['sender'].__name__, 'module_func')
+        self.assertTrue('start' in foo, foo)
+        self.assertGreater(foo['start'], 0)
+        self.assertTrue('stop' not in foo, foo)
+        self.assertTrue('args' in foo, foo)
+        self.assertTrue(len(foo['args']), 1)
+        self.assertEquals(foo['args'][0], 'hi')
+        self.assertTrue('kwargs' in foo, foo)
+        self.assertTrue(len(foo['kwargs']), 1)
+        self.assertTrue('foo' in foo['kwargs'])
+        self.assertEquals(foo['kwargs']['foo'], 'bar')
+    
+        callbacks['before'] = {}
+    
+        @pre_dispatch(TrackingTestCase.class_func)
+        def test(**kwargs):
+            foo.update(kwargs)
+    
+        self.assertTrue(hasattr(TrackingTestCase.class_func, '__wrapped__'))
+        self.assertEquals(len(callbacks['before']), 1)
+
+        self.class_func('hello', foo='bar')
+
+        self.assertTrue('sender' in foo, foo)
+        # best we can do
+        self.assertEquals(foo['sender'].__name__, 'class_func')
+        self.assertTrue('start' in foo, foo)
+        self.assertGreater(foo['start'], 0)
+        self.assertTrue('stop' not in foo, foo)
+        self.assertTrue('args' in foo, foo)
+        self.assertTrue(len(foo['args']), 2)
+        self.assertEquals(foo['args'][1], 'hello')
+        self.assertTrue('kwargs' in foo, foo)
+        self.assertTrue(len(foo['kwargs']), 1)
+        self.assertTrue('foo' in foo['kwargs'])
+        self.assertEquals(foo['kwargs']['foo'], 'bar')
+
+        # callbacks['before'] = {}
+        #     
+        #         @pre_dispatch(TrackingTestCase.class_method)
+        #         def test(**kwargs):
+        #             foo.update(kwargs)
+        #     
+        #         self.assertTrue(hasattr(TrackingTestCase.class_method, '__wrapped__'))
+        #         self.assertEquals(len(callbacks['before']), 1)
+        # 
+        #         TrackingTestCase.class_method()
+        # 
+        #         self.assertTrue('sender' in foo, foo)
+        #         # best we can do
+        #         self.assertEquals(foo['sender'].__name__, 'class_method')
+        #         self.assertTrue('start' in foo, foo)
+        #         self.assertTrue('stop' not in foo, foo)
+        #         self.assertTrue('args' in foo, foo)
+
+    def test_post_hook(self):
+        foo = {}
+        
+        @post_dispatch(module_func)
+        def test(**kwargs):
+            foo.update(kwargs)
+            
+        self.assertTrue(hasattr(module_func, '__wrapped__'))
+        self.assertEquals(len(callbacks['after']), 1)
+        
+        module_func('hi', foo='bar')
+        
+        self.assertTrue('sender' in foo, foo)
+        # best we can do
+        self.assertEquals(foo['sender'].__name__, 'module_func')
+        self.assertTrue('start' in foo, foo)
+        self.assertGreater(foo['start'], 0)
+        self.assertTrue('stop' in foo, foo)
+        self.assertGreater(foo['stop'], foo['start'])
+        self.assertTrue('args' in foo, foo)
+        self.assertTrue(len(foo['args']), 1)
+        self.assertEquals(foo['args'][0], 'hi')
+        self.assertTrue('kwargs' in foo, foo)
+        self.assertTrue(len(foo['kwargs']), 1)
+        self.assertTrue('foo' in foo['kwargs'])
+        self.assertEquals(foo['kwargs']['foo'], 'bar')
+    
+        callbacks['after'] = {}
+    
+        @post_dispatch(TrackingTestCase.class_func)
+        def test(**kwargs):
+            foo.update(kwargs)
+    
+        self.assertTrue(hasattr(TrackingTestCase.class_func, '__wrapped__'))
+        self.assertEquals(len(callbacks['after']), 1)
+
+        self.class_func('hello', foo='bar')
+
+        self.assertTrue('sender' in foo, foo)
+        # best we can do
+        self.assertEquals(foo['sender'].__name__, 'class_func')
+        self.assertTrue('start' in foo, foo)
+        self.assertGreater(foo['start'], 0)
+        self.assertTrue('stop' in foo, foo)
+        self.assertGreater(foo['stop'], foo['start'])
+        self.assertTrue('args' in foo, foo)
+        self.assertTrue(len(foo['args']), 2)
+        self.assertEquals(foo['args'][1], 'hello')
+        self.assertTrue('kwargs' in foo, foo)
+        self.assertTrue(len(foo['kwargs']), 1)
+        self.assertTrue('foo' in foo['kwargs'])
+        self.assertEquals(foo['kwargs']['foo'], 'bar')
\ No newline at end of file