_assert_num_queries: support any connection

Author: blueyed

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
@@ -270,8 +275,11 @@ Example
 ``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 currently only supports the default database.
+
+It is a specialized version of :fixture:`django_assert_num_queries`.
 
 
 Example

pytest_django/fixtures.py

@@ -355,9 +355,12 @@ def _live_server_helper(request):
 
 
 @contextmanager
-def _assert_num_queries(config, num, exact=True):
-    from django.db import connection
+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

tests/test_fixtures.py

@@ -138,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"""