Merge pull request #547 from anrie/feat/max_num_queries

Closes #547.

Author: blueyed

docs/changelog.rst

@@ -1,6 +1,16 @@
 Changelog
 =========
 
+unreleased
+------------------
+
+Features
+^^^^^^^^
+
+* Added new fixture :fixture:`django_assert_max_num_queries` (#547).
+* Assed support for ``connection`` and returning the wrapped context manager
+  with :fixture:`django_assert_num_queries` (#547).
+
 3.3.3 (2018-07-26)
 ------------------
 

docs/helpers.rst

@@ -249,8 +249,13 @@ Example
 ``django_assert_num_queries``
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
+.. fixture:: django_assert_num_queries
+
 This fixture allows to check for an expected number of DB queries.
-It currently only supports the default database.
+
+It wraps `django.test.utils.CaptureQueriesContext`.  A non-default DB
+connection can be passed in using the `connection` keyword argument, and it
+will yield the wrapped CaptureQueriesContext instance.
 
 
 Example
@@ -259,11 +264,34 @@ Example
 ::
 
     def test_queries(django_assert_num_queries):
-        with django_assert_num_queries(3):
+        with django_assert_num_queries(3) as captured:
             Item.objects.create('foo')
             Item.objects.create('bar')
             Item.objects.create('baz')
 
+        assert 'foo' in captured.captured_queries[0]['sql']
+
+
+``django_assert_max_num_queries``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. fixture:: django_assert_max_num_queries
+
+This fixture allows to check for an expected maximum number of DB queries.
+
+It is a specialized version of :fixture:`django_assert_num_queries`.
+
+
+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,39 @@ def _live_server_helper(request):
     request.addfinalizer(live_server._live_server_modified_settings.disable)
 
 
+@contextmanager
+def _assert_num_queries(config, num, exact=True, connection=None):
+    from django.test.utils import CaptureQueriesContext
+
+    if connection is None:
+        from django.db import connection
+
+    verbose = config.getoption('verbose') > 0
+    with CaptureQueriesContext(connection) as context:
+        yield context
+        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 {} done'.format(
+                    num_queries == 1 and '1 was' or '%d were' % (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)
+
+
 @pytest.fixture(scope='function')
 def django_assert_num_queries(pytestconfig):
-    from django.db import connection
-    from django.test.utils import CaptureQueriesContext
+    return partial(_assert_num_queries, pytestconfig)
 
-    @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_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

@@ -59,13 +59,37 @@ def test_django_assert_num_queries_db(django_assert_num_queries):
         Item.objects.create(name='bar')
         Item.objects.create(name='baz')
 
-    with pytest.raises(pytest.fail.Exception):
-        with django_assert_num_queries(2):
+    with pytest.raises(pytest.fail.Exception) as excinfo:
+        with django_assert_num_queries(2) as captured:
             Item.objects.create(name='quux')
+    assert excinfo.value.args == (
+        'Expected to perform 2 queries but 1 was done '
+        '(add -v option to show queries)',)
+    assert len(captured.captured_queries) == 1
+
+
+@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) as excinfo:
+        with django_assert_max_num_queries(2) as captured:
+            Item.objects.create(name='1-foo')
+            Item.objects.create(name='2-bar')
+            Item.objects.create(name='3-quux')
+
+    assert excinfo.value.args == (
+        'Expected to perform 2 queries or less but 3 were done '
+        '(add -v option to show queries)',)
+    assert len(captured.captured_queries) == 3
+    assert '1-foo' in captured.captured_queries[0]['sql']
 
 
 @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):
@@ -114,6 +138,21 @@ def test_queries(django_assert_num_queries):
     assert result.ret == 1
 
 
+@pytest.mark.django_db
+def test_django_assert_num_queries_db_connection(django_assert_num_queries):
+    from django.db import connection
+
+    with django_assert_num_queries(1, connection=connection):
+        Item.objects.create(name='foo')
+
+    with django_assert_num_queries(1, connection=None):
+        Item.objects.create(name='foo')
+
+    with pytest.raises(AttributeError):
+        with django_assert_num_queries(1, connection=False):
+            pass
+
+
 class TestSettings:
     """Tests for the settings fixture, order matters"""