Add django_assert_num_queries fixture (#387)

Author: lukaszb

docs/helpers.rst

@@ -218,6 +218,26 @@ Example
         settings.USE_TZ = True
         assert settings.USE_TZ
 
+
+``django_assert_num_queries``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+This fixture allows to check for an expected number of DB queries.
+It currently only supports the default database.
+
+
+Example
+"""""""
+
+::
+
+    def test_queries(assert_num_queries):
+        with django_assert_num_queries(3):
+            Item.objects.create('foo')
+            Item.objects.create('bar')
+            Item.objects.create('baz')
+
+
 ``mailoutbox``
 ~~~~~~~~~~~~~~~~~~~~~~~~~
 

pytest_django/fixtures.py

@@ -6,6 +6,8 @@
 
 import pytest
 
+from contextlib import contextmanager
+
 from . import live_server_helper
 
 from .django_compat import is_django_unittest
@@ -16,7 +18,7 @@
 __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']
+           '_live_server_helper', 'django_assert_num_queries']
 
 
 @pytest.fixture(scope='session')
@@ -339,3 +341,24 @@ def _live_server_helper(request):
     """
     if 'live_server' in request.funcargnames:
         getfixturevalue(request, 'transactional_db')
+
+
+@pytest.fixture(scope='function')
+def django_assert_num_queries(pytestconfig):
+    from django.db import connection
+    from django.test.utils import CaptureQueriesContext
+
+    @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_django/plugin.py

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

@@ -8,7 +8,7 @@
 
 import pytest
 
-from django.db import connection
+from django.db import connection, transaction
 from django.conf import settings as real_settings
 from django.core import mail
 from django.test.client import Client, RequestFactory
@@ -50,6 +50,68 @@ def test_rf(rf):
     assert isinstance(rf, RequestFactory)
 
 
+@pytest.mark.django_db
+def test_django_assert_num_queries_db(django_assert_num_queries):
+    with django_assert_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):
+            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):
+    with transaction.atomic():
+
+        with django_assert_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):
+                Item.objects.create(name='quux')
+
+
+def test_django_assert_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):
+                list(ContentType.objects.all())
+                ContentType.objects.count()
+    """)
+    result = django_testdir.runpytest_subprocess('--tb=short')
+    result.stdout.fnmatch_lines(['*Expected to perform 1 queries but 2 were done*'])
+    assert result.ret == 1
+
+
+def test_django_assert_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):
+                list(ContentType.objects.all())
+                ContentType.objects.count()
+    """)
+    result = django_testdir.runpytest_subprocess('--tb=short', '-v')
+    result.stdout.fnmatch_lines([
+        '*Expected to perform 11 queries but 2 were done*',
+        '*Queries:*',
+        '*========*',
+    ])
+    assert result.ret == 1
+
+
 class TestSettings:
     """Tests for the settings fixture, order matters"""