Add fixture for asserting maximum number of database queries

Author: anrie

docs/helpers.rst

@@ -265,6 +265,24 @@ Example
             Item.objects.create('baz')
 
 
+``django_assert_max_num_queries``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+This fixture allows to check for an expected maximum number of DB queries.
+It currently only supports the default database.
+
+
+Example
+"""""""
+
+::
+
+    def test_max_queries(django_assert_max_num_queries):
+        with django_assert_max_num_queries(3):
+            Item.objects.create('foo')
+            Item.objects.create('bar')
+
+
 ``mailoutbox``
 ~~~~~~~~~~~~~~
 

pytest_django/fixtures.py

@@ -3,6 +3,7 @@
 from __future__ import with_statement
 
 import os
+from functools import partial
 
 import pytest
 
@@ -17,7 +18,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_num_queries',
+           'django_assert_max_num_queries']
 
 
 @pytest.fixture(scope='session')
@@ -352,22 +354,34 @@ def _live_server_helper(request):
     request.addfinalizer(live_server._live_server_modified_settings.disable)
 
 
-@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
+        num_queries = len(context)
+        failed = num != num_queries if exact else num < num_queries
+        if failed:
+            msg = "Expected to perform {} queries {}{}".format(
+                num,
+                '' if exact else 'or less ',
+                'but {} were done'.format(num_queries)
+            )
+            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_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

@@ -16,6 +16,7 @@
 
 from .django_compat import is_django_unittest  # noqa
 from .fixtures import django_assert_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

@@ -64,8 +64,22 @@ def test_django_assert_num_queries_db(django_assert_num_queries):
             Item.objects.create(name='quux')
 
 
+@pytest.mark.django_db
+def test_django_assert_max_num_queries_db(django_assert_max_num_queries):
+    with django_assert_max_num_queries(2):
+        Item.objects.create(name='1-foo')
+        Item.objects.create(name='2-bar')
+
+    with pytest.raises(pytest.fail.Exception):
+        with django_assert_max_num_queries(2):
+            Item.objects.create(name='1-foo')
+            Item.objects.create(name='2-bar')
+            Item.objects.create(name='3-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_num_queries_transactional_db(
+        transactional_db, django_assert_num_queries):
     with transaction.atomic():
 
         with django_assert_num_queries(3):