Add fixture for asserting maximum number of database queries

Author: anrie

docs/helpers.rst

@@ -220,7 +220,7 @@ Example
         assert settings.USE_TZ
 
 
-``django_assert_num_queries``
+``django_assert_exact_num_queries``
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 This fixture allows to check for an expected number of DB queries.
@@ -232,8 +232,8 @@ Example
 
 ::
 
-    def test_queries(django_assert_num_queries):
-        with django_assert_num_queries(3):
+    def test_queries(django_assert_exact_num_queries):
+        with django_assert_exact_num_queries(3):
             Item.objects.create('foo')
             Item.objects.create('bar')
             Item.objects.create('baz')

pytest_django/fixtures.py

@@ -3,6 +3,7 @@
 from __future__ import with_statement
 
 import os
+from functools import partial
 
 import pytest
 
@@ -18,7 +19,8 @@
 __all__ = ['django_db_setup', 'db', 'transactional_db', 'admin_user',
            'django_user_model', 'django_username_field',
            'client', 'admin_client', 'rf', 'settings', 'live_server',
-           '_live_server_helper', 'django_assert_num_queries']
+           '_live_server_helper', 'django_assert_exact_num_queries',
+           'django_assert_max_num_queries']
 
 
 @pytest.fixture(scope='session')
@@ -336,22 +338,33 @@ def _live_server_helper(request):
         getfixturevalue(request, 'transactional_db')
 
 
-@pytest.fixture(scope='function')
-def django_assert_num_queries(pytestconfig):
+@contextmanager
+def _assert_num_queries(config, num, exact=True):
     from django.db import connection
     from django.test.utils import CaptureQueriesContext
+    verbose = config.getoption('verbose') > 0
+    with CaptureQueriesContext(connection) as context:
+        yield
+        failed = num != len(context) if exact else num < len(context)
+        if failed:
+            msg = "Expected to perform {} queries {}{}".format(
+                num,
+                '' if exact else 'or less ',
+                'but {} were done'.format(len(context))
+            )
+            if verbose:
+                sqls = (q['sql'] for q in context.captured_queries)
+                msg += '\n\nQueries:\n========\n\n%s' % '\n\n'.join(sqls)
+            else:
+                msg += " (add -v option to show queries)"
+            pytest.fail(msg)
+
 
-    @contextmanager
-    def _assert_num_queries(num):
-        with CaptureQueriesContext(connection) as context:
-            yield
-            if num != len(context):
-                msg = "Expected to perform %s queries but %s were done" % (num, len(context))
-                if pytestconfig.getoption('verbose') > 0:
-                    sqls = (q['sql'] for q in context.captured_queries)
-                    msg += '\n\nQueries:\n========\n\n%s' % '\n\n'.join(sqls)
-                else:
-                    msg += " (add -v option to show queries)"
-                pytest.fail(msg)
-
-    return _assert_num_queries
+@pytest.fixture(scope='function')
+def django_assert_exact_num_queries(pytestconfig):
+    return partial(_assert_num_queries, pytestconfig)
+
+
+@pytest.fixture(scope='function')
+def django_assert_max_num_queries(pytestconfig):
+    return partial(_assert_num_queries, pytestconfig, exact=False)

pytest_django/plugin.py

@@ -15,7 +15,8 @@
 import pytest
 
 from .django_compat import is_django_unittest  # noqa
-from .fixtures import django_assert_num_queries  # noqa
+from .fixtures import django_assert_exact_num_queries  # noqa
+from .fixtures import django_assert_max_num_queries  # noqa
 from .fixtures import django_db_setup  # noqa
 from .fixtures import django_db_use_migrations  # noqa
 from .fixtures import django_db_keepdb  # noqa

tests/test_fixtures.py

@@ -51,39 +51,40 @@ def test_rf(rf):
 
 
 @pytest.mark.django_db
-def test_django_assert_num_queries_db(django_assert_num_queries):
-    with django_assert_num_queries(3):
+def test_django_assert_exact_num_queries_db(django_assert_exact_num_queries):
+    with django_assert_exact_num_queries(3):
         Item.objects.create(name='foo')
         Item.objects.create(name='bar')
         Item.objects.create(name='baz')
 
     with pytest.raises(pytest.fail.Exception):
-        with django_assert_num_queries(2):
+        with django_assert_exact_num_queries(2):
             Item.objects.create(name='quux')
 
 
 @pytest.mark.django_db(transaction=True)
-def test_django_assert_num_queries_transactional_db(transactional_db, django_assert_num_queries):
+def test_django_assert_exact_num_queries_transactional_db(
+        transactional_db, django_assert_exact_num_queries):
     with transaction.atomic():
 
-        with django_assert_num_queries(3):
+        with django_assert_exact_num_queries(3):
             Item.objects.create(name='foo')
             Item.objects.create(name='bar')
             Item.objects.create(name='baz')
 
         with pytest.raises(pytest.fail.Exception):
-            with django_assert_num_queries(2):
+            with django_assert_exact_num_queries(2):
                 Item.objects.create(name='quux')
 
 
-def test_django_assert_num_queries_output(django_testdir):
+def test_django_assert_exact_num_queries_output(django_testdir):
     django_testdir.create_test_module("""
         from django.contrib.contenttypes.models import ContentType
         import pytest
 
         @pytest.mark.django_db
-        def test_queries(django_assert_num_queries):
-            with django_assert_num_queries(1):
+        def test_queries(django_assert_exact_num_queries):
+            with django_assert_exact_num_queries(1):
                 list(ContentType.objects.all())
                 ContentType.objects.count()
     """)
@@ -92,14 +93,14 @@ def test_queries(django_assert_num_queries):
     assert result.ret == 1
 
 
-def test_django_assert_num_queries_output_verbose(django_testdir):
+def test_django_assert_exact_num_queries_output_verbose(django_testdir):
     django_testdir.create_test_module("""
         from django.contrib.contenttypes.models import ContentType
         import pytest
 
         @pytest.mark.django_db
-        def test_queries(django_assert_num_queries):
-            with django_assert_num_queries(11):
+        def test_queries(django_assert_exact_num_queries):
+            with django_assert_exact_num_queries(11):
                 list(ContentType.objects.all())
                 ContentType.objects.count()
     """)