Ticket #14774: 14774.diff

django/test/testcases.py

@@ -6 +6 @@
 from django.conf import settings
 from django.core import mail
 from django.core.management import call_command
+from django.core.signals import request_started
 from django.core.urlresolvers import clear_url_caches
-from django.db import transaction, connection, connections, DEFAULT_DB_ALIAS
+from django.db import (transaction, connection, connections, DEFAULT_DB_ALIAS, 
+    reset_queries)
 from django.http import QueryDict
 from django.test import _doctest as doctest
 from django.test.client import Client
@@ -215 +217 @@
         self.test_case = test_case
         self.num = num
         self.connection = connection
+        self.executed = 0
+        
+    def reset_queries_wrapper(self, **kwargs):
+        self.executed += len(self.connection.queries) - self.starting_queries
+        self.starting_queries = 0
+        reset_queries(**kwargs)
 
     def __enter__(self):
         self.old_debug_cursor = self.connection.use_debug_cursor
         self.connection.use_debug_cursor = True
         self.starting_queries = len(self.connection.queries)
+        request_started.disconnect(reset_queries)
+        request_started.connect(self.reset_queries_wrapper, weak=False)
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
         self.connection.use_debug_cursor = self.old_debug_cursor
+        request_started.connect(reset_queries)
+        request_started.disconnect(self.reset_queries_wrapper)
         if exc_type is not None:
             return
 
         final_queries = len(self.connection.queries)
-        executed = final_queries - self.starting_queries
-
+        executed = final_queries - self.starting_queries + self.executed
+        
         self.test_case.assertEqual(
             executed, self.num, "%d queries executed, %d expected" % (
                 executed, self.num

tests/regressiontests/test_utils/tests.py

@@ -2 +2 @@
 
 from django.test import TestCase, skipUnlessDBFeature, skipIfDBFeature
 
+from models import Person
 
 if sys.version_info >= (2, 5):
     from tests_25 import AssertNumQueriesTests
@@ -15 +16 @@
         self.assertRaises(ValueError,
             self.assertNumQueries, 2, test_func
         )
+        
+    def test_assert_num_queries_with_client(self):
+        person = Person.objects.create(name='test')
+        
+        self.assertNumQueries(
+            1,
+            self.client.get,
+            '/test_utils/get_person/%s/' % person.pk
+        )
+        
+        self.assertNumQueries(
+            1,
+            self.client.get,
+            '/test_utils/get_person/%s/' % person.pk
+        )
+        
+        def test_func():
+            self.client.get('/test_utils/get_person/%s/' % person.pk)
+            self.client.get('/test_utils/get_person/%s/' % person.pk)
+        self.assertNumQueries(2, test_func)
+            
 
     def test_skip_unless_db_feature(self):
         "A test that might be skipped is actually called."

tests/regressiontests/test_utils/tests_25.py

@@ -26 +26 @@
         with self.assertRaises(TypeError):
             with self.assertNumQueries(4000):
                 raise TypeError
+    
+    def test_with_client(self):
+        person = Person.objects.create(name='test')
+        
+        with self.assertNumQueries(1):
+            self.client.get('/test_utils/get_person/%s/' % person.pk)
+            
+        with self.assertNumQueries(1):
+            self.client.get('/test_utils/get_person/%s/' % person.pk)
+        
+        with self.assertNumQueries(2):
+            self.client.get('/test_utils/get_person/%s/' % person.pk)
+            self.client.get('/test_utils/get_person/%s/' % person.pk)
+ No newline at end of file

new file tests/regressiontests/test_utils/urls.py

@@ +1 @@
+from django.conf.urls.defaults import *
+import views
+
+urlpatterns = patterns('',
+    (r'^get_person/(\d+)/$', views.get_person),
+)
+ No newline at end of file

new file tests/regressiontests/test_utils/views.py

@@ +1 @@
+from django.http import HttpResponse
+from django.shortcuts import get_object_or_404
+from models import Person
+
+def get_person(request, pk):
+    person = get_object_or_404(Person, pk=pk)
+    return HttpResponse(person.name)
+ No newline at end of file

tests/urls.py

@@ -41 +41 @@
 
     # special headers views
     (r'special_headers/', include('regressiontests.special_headers.urls')),
+    
+    # test util views
+    (r'test_utils/', include('regressiontests.test_utils.urls')),
 )