33 lines
921 B
Python
33 lines
921 B
Python
|
from django.db import DEFAULT_DB_ALIAS, connections
|
||
|
|
from django.test.utils import CaptureQueriesContext
|
||
|
|
|
||
|
|
|
||
|
|
def count_queries(func, verbose=False, *args, **kwargs) -> int:
|
||
|
|
context = CaptureQueriesContext(connections[DEFAULT_DB_ALIAS])
|
||
|
|
with context:
|
||
|
|
func(*args, **kwargs)
|
||
|
|
|
||
|
|
if verbose:
|
||
|
|
queries = "\n".join(
|
||
|
|
f"{i}. {query['sql']}"
|
||
|
|
for i, query in enumerate(context.captured_queries, start=1)
|
||
|
|
)
|
||
|
|
print(f"{len(context)} queries executed\nCaptured queries were:\n{queries}")
|
||
|
|
|
||
|
|
return len(context)
|
||
|
|
|
||
|
|
|
||
|
|
def assert_query_count(should_be, verbose=False):
|
||
|
|
"""
|
||
|
|
Decorator to easily count queries on any test you want.
|
||
|
|
should_be defines how many queries are to be expected
|
||
|
|
"""
|
||
|
|
|
||
|
|
def outer(func):
|
||
|
|
def inner(*args, **kwargs):
|
||
|
|
assert count_queries(func, verbose, *args, **kwargs) == should_be
|
||
|
|
|
||
|
|
return inner
|
||
|
|
|
||
|
|
return outer
|