OpenSlides/tests/count_queries.py
2020-03-17 07:24:36 +01:00

47 lines
1.4 KiB
Python

from typing import Callable
from django.db import DEFAULT_DB_ALIAS, connections
from django.test.utils import CaptureQueriesContext
def count_queries(func, verbose=False) -> Callable[..., int]:
def wrapper(*args, **kwargs) -> int:
context = CaptureQueriesContext(connections[DEFAULT_DB_ALIAS])
with context:
func(*args, **kwargs)
if verbose:
print(get_verbose_queries(context))
return len(context)
return wrapper
class AssertNumQueriesContext(CaptureQueriesContext):
def __init__(self, test_case, num, verbose):
self.test_case = test_case
self.num = num
self.verbose = verbose
super().__init__(connections[DEFAULT_DB_ALIAS])
def __exit__(self, exc_type, exc_value, traceback):
super().__exit__(exc_type, exc_value, traceback)
if exc_type is not None:
return
executed = len(self)
verbose_queries = get_verbose_queries(self)
if self.verbose:
print(verbose_queries)
self.test_case.assertEqual(executed, self.num)
else:
self.test_case.assertEqual(executed, self.num, verbose_queries)
def get_verbose_queries(context):
queries = "\n".join(
f"{i}. {query['sql']}"
for i, query in enumerate(context.captured_queries, start=1)
)
return f"{len(context)} queries executed\nCaptured queries were:\n{queries}"