Asserting query counts in Django tests with
CaptureQueriesContext is a great way to maintain/optimize code, but it doesn't provide enough context to find what changed when tests fail during a refactor. I found it helpful to capture some lines from the stack trace when each query is being executed and have them accessible during test runs that use
CaptureQueriesContext to be logged along with the query, when count assertions fail.
You can achieve this by replacing
django.db.backends.utils with a new class. The new class can inherit from
CursorDebugWrapper and override
executemany methods. The replacement methods would wrap the call (of
super()) with a
try/finally block, and in the
finally block, call a new private method that can collect the stack trace.
Real world example
The following example assumes you have an environment variable
ENABLE_QUERY_TRACEBACK_CAPTURE=1 set during test runs to override
Somewhere in the project, say a
import traceback from typing import List from dataclasses import dataclass from django.db.backends.utils import CursorDebugWrapper from django.conf import settings @dataclass(frozen=True) class Traceline: path: str lineno: int func: str line: str @dataclass(frozen=True) class QueryTrace: sql: str traceback: List[Traceline] class CustomCursorDebugWrapper(CursorDebugWrapper): def __init__(self, cursor, db): super().__init__(cursor, db) if not hasattr(self.db, "captured_query_traceback"): self.db.captured_query_traceback =  def capture_traceback(self, sql): stack = traceback.extract_stack() self.db.captured_query_traceback.append( QueryTrace( sql=sql, traceback=[ Traceline( path=path, lineno=lineno, func=func, line=line, ) for (path, lineno, func, line) in reversed(stack) ], ) ) def execute(self, sql, params=None): try: return super().execute(sql, params) finally: sql = self.db.ops.last_executed_query(self.cursor, sql, params) self.capture_traceback(sql) def executemany(self, sql, param_list): try: return super().executemany(sql, param_list) finally: self.capture_traceback(sql)
import os if os.environ.get("ENABLE_QUERY_TRACEBACK_CAPTURE"): import django.db.backends.utils as backend_utils from testutils import CustomCursorDebugWrapper backend_utils.CursorDebugWrapper = CustomCursorDebugWrapper