Skip to content

Commit

Permalink
Adds using to django_assert_num_queries (#1170)
Browse files Browse the repository at this point in the history
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...
  • Loading branch information
kingbuzzman authored Feb 10, 2025
1 parent c3018d6 commit 0ee43ef
Show file tree
Hide file tree
Showing 4 changed files with 46 additions and 7 deletions.
4 changes: 4 additions & 0 deletions docs/changelog.rst
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,10 @@ Compatibility

* Added official support for Python 3.13.

* Added ``using`` argument to :fixture:`django_assert_num_queries` and
:fixture:`django_assert_max_num_queries` to easily specify the database
alias to use.

Bugfixes
^^^^^^^^

Expand Down
10 changes: 6 additions & 4 deletions docs/helpers.rst
Original file line number Diff line number Diff line change
Expand Up @@ -423,11 +423,12 @@ Example
``django_assert_num_queries``
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

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

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

Expand Down Expand Up @@ -462,11 +463,12 @@ If you use type annotations, you can annotate the fixture like this::
``django_assert_max_num_queries``
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

.. py:function:: django_assert_max_num_queries(num, connection=None, info=None)
.. py:function:: django_assert_max_num_queries(num, connection=None, info=None, *, using=None)
:param num: expected maximum number of queries
:param connection: optional non-default DB connection
:param connection: optional database connection
:param str info: optional info message to display on failure
:param str using: optional database alias

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

Expand Down
16 changes: 13 additions & 3 deletions pytest_django/fixtures.py
Original file line number Diff line number Diff line change
Expand Up @@ -606,6 +606,8 @@ def __call__(
num: int,
connection: Any | None = ...,
info: str | None = ...,
*,
using: str | None = ...,
) -> django.test.utils.CaptureQueriesContext:
pass # pragma: no cover

Expand All @@ -617,13 +619,21 @@ def _assert_num_queries(
exact: bool = True,
connection: Any | None = None,
info: str | None = None,
*,
using: str | None = None,
) -> Generator[django.test.utils.CaptureQueriesContext, None, None]:
from django.db import connection as default_conn, connections
from django.test.utils import CaptureQueriesContext

if connection is None:
from django.db import connection as conn
else:
if connection and using:
raise ValueError('The "connection" and "using" parameter cannot be used together')

if connection is not None:
conn = connection
elif using is not None:
conn = connections[using]
else:
conn = default_conn

verbose = config.getoption("verbose") > 0
with CaptureQueriesContext(conn) as context:
Expand Down
23 changes: 23 additions & 0 deletions tests/test_fixtures.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
from django.core import mail
from django.db import connection, transaction
from django.test import AsyncClient, AsyncRequestFactory, Client, RequestFactory
from django.utils.connection import ConnectionDoesNotExist
from django.utils.encoding import force_str

from .helpers import DjangoPytester
Expand Down Expand Up @@ -206,6 +207,28 @@ def test_django_assert_num_queries_db_connection(
pass


@pytest.mark.django_db
def test_django_assert_num_queries_db_using(
django_assert_num_queries: DjangoAssertNumQueries,
) -> None:
from django.db import connection

with django_assert_num_queries(1, using="default"):
Item.objects.create(name="foo")

error_message = 'The "connection" and "using" parameter cannot be used together'
with pytest.raises(ValueError, match=error_message):
with django_assert_num_queries(1, connection=connection, using="default"):
Item.objects.create(name="foo")

with django_assert_num_queries(1, using=None):
Item.objects.create(name="foo")

with pytest.raises(ConnectionDoesNotExist):
with django_assert_num_queries(1, using="bad_db_name"):
pass


@pytest.mark.django_db
def test_django_assert_num_queries_output_info(django_pytester: DjangoPytester) -> None:
django_pytester.create_test_module(
Expand Down

0 comments on commit 0ee43ef

Please sign in to comment.