Skip to content

Commit 0ee43ef

Browse files
authored
Adds using to django_assert_num_queries (#1170)
Adds quality of life to replace ```py from django.db import connections def test_test(): with django_assert_num_queries(1, connection=connections["log"]): Model.objects.get() ``` with ```py def test_test(): with django_assert_num_queries(1, using="log"): Model.objects.get() ``` I have multiple DBs and this is driving me nuts...
1 parent c3018d6 commit 0ee43ef

File tree

4 files changed

+46
-7
lines changed

4 files changed

+46
-7
lines changed

docs/changelog.rst

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,10 @@ Compatibility
99

1010
* Added official support for Python 3.13.
1111

12+
* Added ``using`` argument to :fixture:`django_assert_num_queries` and
13+
:fixture:`django_assert_max_num_queries` to easily specify the database
14+
alias to use.
15+
1216
Bugfixes
1317
^^^^^^^^
1418

docs/helpers.rst

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -423,11 +423,12 @@ Example
423423
``django_assert_num_queries``
424424
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
425425

426-
.. py:function:: django_assert_num_queries(num, connection=None, info=None)
426+
.. py:function:: django_assert_num_queries(num, connection=None, info=None, *, using=None)
427427
428428
:param num: expected number of queries
429-
:param connection: optional non-default DB connection
429+
:param connection: optional database connection
430430
:param str info: optional info message to display on failure
431+
:param str using: optional database alias
431432

432433
This fixture allows to check for an expected number of DB queries.
433434

@@ -462,11 +463,12 @@ If you use type annotations, you can annotate the fixture like this::
462463
``django_assert_max_num_queries``
463464
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
464465

465-
.. py:function:: django_assert_max_num_queries(num, connection=None, info=None)
466+
.. py:function:: django_assert_max_num_queries(num, connection=None, info=None, *, using=None)
466467
467468
:param num: expected maximum number of queries
468-
:param connection: optional non-default DB connection
469+
:param connection: optional database connection
469470
:param str info: optional info message to display on failure
471+
:param str using: optional database alias
470472

471473
This fixture allows to check for an expected maximum number of DB queries.
472474

pytest_django/fixtures.py

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -606,6 +606,8 @@ def __call__(
606606
num: int,
607607
connection: Any | None = ...,
608608
info: str | None = ...,
609+
*,
610+
using: str | None = ...,
609611
) -> django.test.utils.CaptureQueriesContext:
610612
pass # pragma: no cover
611613

@@ -617,13 +619,21 @@ def _assert_num_queries(
617619
exact: bool = True,
618620
connection: Any | None = None,
619621
info: str | None = None,
622+
*,
623+
using: str | None = None,
620624
) -> Generator[django.test.utils.CaptureQueriesContext, None, None]:
625+
from django.db import connection as default_conn, connections
621626
from django.test.utils import CaptureQueriesContext
622627

623-
if connection is None:
624-
from django.db import connection as conn
625-
else:
628+
if connection and using:
629+
raise ValueError('The "connection" and "using" parameter cannot be used together')
630+
631+
if connection is not None:
626632
conn = connection
633+
elif using is not None:
634+
conn = connections[using]
635+
else:
636+
conn = default_conn
627637

628638
verbose = config.getoption("verbose") > 0
629639
with CaptureQueriesContext(conn) as context:

tests/test_fixtures.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
from django.core import mail
1616
from django.db import connection, transaction
1717
from django.test import AsyncClient, AsyncRequestFactory, Client, RequestFactory
18+
from django.utils.connection import ConnectionDoesNotExist
1819
from django.utils.encoding import force_str
1920

2021
from .helpers import DjangoPytester
@@ -206,6 +207,28 @@ def test_django_assert_num_queries_db_connection(
206207
pass
207208

208209

210+
@pytest.mark.django_db
211+
def test_django_assert_num_queries_db_using(
212+
django_assert_num_queries: DjangoAssertNumQueries,
213+
) -> None:
214+
from django.db import connection
215+
216+
with django_assert_num_queries(1, using="default"):
217+
Item.objects.create(name="foo")
218+
219+
error_message = 'The "connection" and "using" parameter cannot be used together'
220+
with pytest.raises(ValueError, match=error_message):
221+
with django_assert_num_queries(1, connection=connection, using="default"):
222+
Item.objects.create(name="foo")
223+
224+
with django_assert_num_queries(1, using=None):
225+
Item.objects.create(name="foo")
226+
227+
with pytest.raises(ConnectionDoesNotExist):
228+
with django_assert_num_queries(1, using="bad_db_name"):
229+
pass
230+
231+
209232
@pytest.mark.django_db
210233
def test_django_assert_num_queries_output_info(django_pytester: DjangoPytester) -> None:
211234
django_pytester.create_test_module(

0 commit comments

Comments
 (0)