Fixed #5416 -- Added TestCase.assertNumQueries, which tests that a given function executes the correct number of queries.

git-svn-id: http://code.djangoproject.com/svn/django/trunk@14183 bcc190cf-cafb-0310-a4f2-bffc1f526a37

7 files changed, 100 insertions, 93 deletions

diff --git a/tests/regressiontests/defer_regress/tests.py b/tests/regressiontests/defer_regress/tests.py
index affb0e2405..6f1b023b5b 100644
--- a/tests/regressiontests/defer_regress/tests.py
+++ b/tests/regressiontests/defer_regress/tests.py
@@ -11,17 +11,6 @@ from models import ResolveThis, Item, RelatedItem, Child, Leaf
 
 
 class DeferRegressionTest(TestCase):
-    def assert_num_queries(self, n, func, *args, **kwargs):
-        old_DEBUG = settings.DEBUG
-        settings.DEBUG = True
-        starting_queries = len(connection.queries)
-        try:
-            func(*args, **kwargs)
-        finally:
-            settings.DEBUG = old_DEBUG
-        self.assertEqual(starting_queries + n, len(connection.queries))
-
-
     def test_basic(self):
         # Deferred fields should really be deferred and not accidentally use
         # the field's default value just because they aren't passed to __init__
@@ -33,19 +22,19 @@ class DeferRegressionTest(TestCase):
         def test():
             self.assertEqual(obj.name, "first")
             self.assertEqual(obj.other_value, 0)
-        self.assert_num_queries(0, test)
+        self.assertNumQueries(0, test)
 
         def test():
             self.assertEqual(obj.value, 42)
-        self.assert_num_queries(1, test)
+        self.assertNumQueries(1, test)
 
         def test():
             self.assertEqual(obj.text, "xyzzy")
-        self.assert_num_queries(1, test)
+        self.assertNumQueries(1, test)
 
         def test():
             self.assertEqual(obj.text, "xyzzy")
-        self.assert_num_queries(0, test)
+        self.assertNumQueries(0, test)
 
         # Regression test for #10695. Make sure different instances don't
         # inadvertently share data in the deferred descriptor objects.
diff --git a/tests/regressiontests/forms/models.py b/tests/regressiontests/forms/models.py
index 028ff9bad2..a4891df06e 100644
--- a/tests/regressiontests/forms/models.py
+++ b/tests/regressiontests/forms/models.py
@@ -1,10 +1,9 @@
 # -*- coding: utf-8 -*-
 import datetime
-import tempfile
 import shutil
+import tempfile
 
-from django.db import models, connection
-from django.conf import settings
+from django.db import models
 # Can't import as "forms" due to implementation details in the test suite (the
 # current file is called "forms" and is already imported).
 from django import forms as django_forms
@@ -77,19 +76,13 @@ class TestTicket12510(TestCase):
     ''' It is not necessary to generate choices for ModelChoiceField (regression test for #12510). '''
     def setUp(self):
         self.groups = [Group.objects.create(name=name) for name in 'abc']
-        self.old_debug = settings.DEBUG
-        # turn debug on to get access to connection.queries
-        settings.DEBUG = True
-
-    def tearDown(self):
-        settings.DEBUG = self.old_debug
 
     def test_choices_not_fetched_when_not_rendering(self):
-        initial_queries = len(connection.queries)
-        field = django_forms.ModelChoiceField(Group.objects.order_by('-name'))
-        self.assertEqual('a', field.clean(self.groups[0].pk).name)
+        def test():
+            field = django_forms.ModelChoiceField(Group.objects.order_by('-name'))
+            self.assertEqual('a', field.clean(self.groups[0].pk).name)
         # only one query is required to pull the model from DB
-        self.assertEqual(initial_queries+1, len(connection.queries))
+        self.assertNumQueries(1, test)
 
 class ModelFormCallableModelDefault(TestCase):
     def test_no_empty_option(self):
diff --git a/tests/regressiontests/model_forms_regress/tests.py b/tests/regressiontests/model_forms_regress/tests.py
index 397651a6b8..d695104d50 100644
--- a/tests/regressiontests/model_forms_regress/tests.py
+++ b/tests/regressiontests/model_forms_regress/tests.py
@@ -1,10 +1,8 @@
 import unittest
 from datetime import date
 
-from django import db
 from django import forms
 from django.forms.models import modelform_factory, ModelChoiceField
-from django.conf import settings
 from django.test import TestCase
 from django.core.exceptions import FieldError, ValidationError
 from django.core.files.uploadedfile import SimpleUploadedFile
@@ -14,14 +12,6 @@ from models import Person, RealPerson, Triple, FilePathModel, Article, \
 
 
 class ModelMultipleChoiceFieldTests(TestCase):
-
-    def setUp(self):
-        self.old_debug = settings.DEBUG
-        settings.DEBUG = True
-
-    def tearDown(self):
-        settings.DEBUG = self.old_debug
-
     def test_model_multiple_choice_number_of_queries(self):
         """
         Test that ModelMultipleChoiceField does O(1) queries instead of
@@ -30,10 +20,8 @@ class ModelMultipleChoiceFieldTests(TestCase):
         for i in range(30):
             Person.objects.create(name="Person %s" % i)
 
-        db.reset_queries()
         f = forms.ModelMultipleChoiceField(queryset=Person.objects.all())
-        selected = f.clean([1, 3, 5, 7, 9])
-        self.assertEquals(len(db.connection.queries), 1)
+        self.assertNumQueries(1, f.clean, [1, 3, 5, 7, 9])
 
 class TripleForm(forms.ModelForm):
     class Meta:
@@ -312,7 +300,7 @@ class InvalidFieldAndFactory(TestCase):
                     model = Person
                     fields = ('name', 'no-field')
         except FieldError, e:
-            # Make sure the exception contains some reference to the 
+            # Make sure the exception contains some reference to the
             # field responsible for the problem.
             self.assertTrue('no-field' in e.args[0])
         else:
diff --git a/tests/regressiontests/select_related_onetoone/tests.py b/tests/regressiontests/select_related_onetoone/tests.py
index 4ccb58440a..ab35feccf1 100644
--- a/tests/regressiontests/select_related_onetoone/tests.py
+++ b/tests/regressiontests/select_related_onetoone/tests.py
@@ -7,11 +7,6 @@ from models import (User, UserProfile, UserStat, UserStatResult, StatDetails,
 
 class ReverseSelectRelatedTestCase(TestCase):
     def setUp(self):
-        # Explicitly enable debug for these tests - we need to count
-        # the queries that have been issued.
-        self.old_debug = settings.DEBUG
-        settings.DEBUG = True
-
         user = User.objects.create(username="test")
         userprofile = UserProfile.objects.create(user=user, state="KS",
                                                  city="Lawrence")
@@ -26,65 +21,66 @@ class ReverseSelectRelatedTestCase(TestCase):
                                                   results=results2)
         StatDetails.objects.create(base_stats=advstat, comments=250)
 
-        db.reset_queries()
-
-    def assertQueries(self, queries):
-        self.assertEqual(len(db.connection.queries), queries)
-
-    def tearDown(self):
-        settings.DEBUG = self.old_debug
-
     def test_basic(self):
-        u = User.objects.select_related("userprofile").get(username="test")
-        self.assertEqual(u.userprofile.state, "KS")
-        self.assertQueries(1)
+        def test():
+            u = User.objects.select_related("userprofile").get(username="test")
+            self.assertEqual(u.userprofile.state, "KS")
+        self.assertNumQueries(1, test)
 
     def test_follow_next_level(self):
-        u = User.objects.select_related("userstat__results").get(username="test")
-        self.assertEqual(u.userstat.posts, 150)
-        self.assertEqual(u.userstat.results.results, 'first results')
-        self.assertQueries(1)
+        def test():
+            u = User.objects.select_related("userstat__results").get(username="test")
+            self.assertEqual(u.userstat.posts, 150)
+            self.assertEqual(u.userstat.results.results, 'first results')
+        self.assertNumQueries(1, test)
 
     def test_follow_two(self):
-        u = User.objects.select_related("userprofile", "userstat").get(username="test")
-        self.assertEqual(u.userprofile.state, "KS")
-        self.assertEqual(u.userstat.posts, 150)
-        self.assertQueries(1)
+        def test():
+            u = User.objects.select_related("userprofile", "userstat").get(username="test")
+            self.assertEqual(u.userprofile.state, "KS")
+            self.assertEqual(u.userstat.posts, 150)
+        self.assertNumQueries(1, test)
 
     def test_follow_two_next_level(self):
-        u = User.objects.select_related("userstat__results", "userstat__statdetails").get(username="test")
-        self.assertEqual(u.userstat.results.results, 'first results')
-        self.assertEqual(u.userstat.statdetails.comments, 259)
-        self.assertQueries(1)
+        def test():
+            u = User.objects.select_related("userstat__results", "userstat__statdetails").get(username="test")
+            self.assertEqual(u.userstat.results.results, 'first results')
+            self.assertEqual(u.userstat.statdetails.comments, 259)
+        self.assertNumQueries(1, test)
 
     def test_forward_and_back(self):
-        stat = UserStat.objects.select_related("user__userprofile").get(user__username="test")
-        self.assertEqual(stat.user.userprofile.state, 'KS')
-        self.assertEqual(stat.user.userstat.posts, 150)
-        self.assertQueries(1)
+        def test():
+            stat = UserStat.objects.select_related("user__userprofile").get(user__username="test")
+            self.assertEqual(stat.user.userprofile.state, 'KS')
+            self.assertEqual(stat.user.userstat.posts, 150)
+        self.assertNumQueries(1, test)
 
     def test_back_and_forward(self):
-        u = User.objects.select_related("userstat").get(username="test")
-        self.assertEqual(u.userstat.user.username, 'test')
-        self.assertQueries(1)
+        def test():
+            u = User.objects.select_related("userstat").get(username="test")
+            self.assertEqual(u.userstat.user.username, 'test')
+        self.assertNumQueries(1, test)
 
     def test_not_followed_by_default(self):
-        u = User.objects.select_related().get(username="test")
-        self.assertEqual(u.userstat.posts, 150)
-        self.assertQueries(2)
+        def test():
+            u = User.objects.select_related().get(username="test")
+            self.assertEqual(u.userstat.posts, 150)
+        self.assertNumQueries(2, test)
 
     def test_follow_from_child_class(self):
-        stat = AdvancedUserStat.objects.select_related('user', 'statdetails').get(posts=200)
-        self.assertEqual(stat.statdetails.comments, 250)
-        self.assertEqual(stat.user.username, 'bob')
-        self.assertQueries(1)
+        def test():
+            stat = AdvancedUserStat.objects.select_related('user', 'statdetails').get(posts=200)
+            self.assertEqual(stat.statdetails.comments, 250)
+            self.assertEqual(stat.user.username, 'bob')
+        self.assertNumQueries(1, test)
 
     def test_follow_inheritance(self):
-        stat = UserStat.objects.select_related('user', 'advanceduserstat').get(posts=200)
-        self.assertEqual(stat.advanceduserstat.posts, 200)
-        self.assertEqual(stat.user.username, 'bob')
-        self.assertEqual(stat.advanceduserstat.user.username, 'bob')
-        self.assertQueries(1)
+        def test():
+            stat = UserStat.objects.select_related('user', 'advanceduserstat').get(posts=200)
+            self.assertEqual(stat.advanceduserstat.posts, 200)
+            self.assertEqual(stat.user.username, 'bob')
+            self.assertEqual(stat.advanceduserstat.user.username, 'bob')
+        self.assertNumQueries(1, test)
 
     def test_nullable_relation(self):
         im = Image.objects.create(name="imag1")
diff --git a/tests/regressiontests/test_utils/models.py b/tests/regressiontests/test_utils/models.py
index e69de29bb2..4da7a07bbf 100644
--- a/tests/regressiontests/test_utils/models.py
+++ b/tests/regressiontests/test_utils/models.py
@@ -0,0 +1,5 @@
+from django.db import models
+
+
+class Person(models.Model):
+    name = models.CharField(max_length=100)
diff --git a/tests/regressiontests/test_utils/python_25.py b/tests/regressiontests/test_utils/python_25.py
new file mode 100644
index 0000000000..a1e8a94d1e
--- /dev/null
+++ b/tests/regressiontests/test_utils/python_25.py
@@ -0,0 +1,30 @@
+from __future__ import with_statement
+
+from django.test import TestCase
+
+from models import Person
+
+
+class AssertNumQueriesTests(TestCase):
+    def test_simple(self):
+        with self.assertNumQueries(0):
+            pass
+
+        with self.assertNumQueries(1):
+            # Guy who wrote Linux
+            Person.objects.create(name="Linus Torvalds")
+
+        with self.assertNumQueries(2):
+            # Guy who owns the bagel place I like
+            Person.objects.create(name="Uncle Ricky")
+            self.assertEqual(Person.objects.count(), 2)
+
+    def test_failure(self):
+        with self.assertRaises(AssertionError) as exc_info:
+            with self.assertNumQueries(2):
+                Person.objects.count()
+        self.assertEqual(str(exc_info.exception), "1 != 2 : 1 queries executed, 2 expected")
+
+        with self.assertRaises(TypeError):
+            with self.assertNumQueries(4000):
+                raise TypeError
diff --git a/tests/regressiontests/test_utils/tests.py b/tests/regressiontests/test_utils/tests.py
index a2539bf8c6..4f92a402cc 100644
--- a/tests/regressiontests/test_utils/tests.py
+++ b/tests/regressiontests/test_utils/tests.py
@@ -1,6 +1,12 @@
-r"""
+import sys
+
+if sys.version_info >= (2, 5):
+    from python_25 import AssertNumQueriesTests
+
+
+__test__ = {"API_TEST": r"""
 # Some checks of the doctest output normalizer.
-# Standard doctests do fairly 
+# Standard doctests do fairly
 >>> from django.utils import simplejson
 >>> from django.utils.xmlutils import SimplerXMLGenerator
 >>> from StringIO import StringIO
@@ -55,7 +61,7 @@ r"""
 >>> produce_json()
 '["foo", {"whiz": 42, "bar": ["baz", null, 1.0, 2]}]'
 
-# XML output is normalized for attribute order, so it doesn't matter 
+# XML output is normalized for attribute order, so it doesn't matter
 # which order XML element attributes are listed in output
 >>> produce_xml()
 '<?xml version="1.0" encoding="UTF-8"?>\n<foo aaa="1.0" bbb="2.0"><bar ccc="3.0">Hello</bar><whiz>Goodbye</whiz></foo>'
@@ -69,4 +75,4 @@ r"""
 >>> produce_xml_fragment()
 '<foo bbb="2.0" aaa="1.0">Hello</foo><bar ddd="4.0" ccc="3.0"></bar>'
 
-"""
\ No newline at end of file
+"""}