Django: Writing Robust Tests for N+1 Problems
TLDR;
For those in a hurry (or more realistically for future me who will want to find this again in 6 months), here's the important bit:
class BookListViewTestCase(TestCase):
def test_view_query_count(self):
"""
Make sure that listing books always triggers the same number of database
queries no matter how many books there are.
"""
Book.objects.create(
author=Author.objects.create(name="test"),
title="test",
)
with CaptureQueriesContext(connection) as ctx:
self.client.get("/books/")
baseline = len(ctx.captured_queries)
for i in range(10):
# Here a tool like factory-boy or model-bakery would really help...
Book.objects.create(
author=Author.objects.create(name=f"test{i}"),
title="test{i}",
)
self.assertNumQueries(baseline):
self.client.get("/books/")And for anyone else, keep reading for some more in-depth explanations (and thanks for tagging along, I hope you enjoy the ride 🚵🏻).
What's the N+1 problem?
The term "N+1 problem" refers to an issue that can arise when using an ORM where you end up making
more database queries than you should. Typically it happens when you fetch related objects inside a loop,
resulting in 1 query for the loop itself, then 1 query per item in the loop. If the loop has
N items, that means N+1 queries in total, hence the name.
To show a practical example, let's say we have the following models (very imaginative, I know):
# models.py
from django.db import models
class Author(models.Model):
name = models.CharField()
class Book(models.Model):
author = models.ForeignKey(Author)
title = models.CharField()Let's also create a simplistic view that returns a list of all books, alongside their author's name:
# views.py
from django.http import JsonResponse
from .models import Book
def list_books(request):
data = {
"books": [
{
"id": book.pk,
"title": book.title,
"author": book.author.name,
}
for book in Book.objects.all()
]
}
return JsonResponse(data)Can you spot where the problem is?
Testing with assertNumQueries
When it comes to counting database queries in a test, Django helpfully ships with a custom assertion metho
called
assertNumQueries
(which you get for free when inheriting from django.test.TestCase).
You use it as a
context manager (meaning with a
with block), like so:
# tests.py
from django.test import TestCase
from .models import Author, Book
class BookListViewTestCase(TestCase):
def test_view_query_count(self):
"""
Make sure that listing books always triggers the same number of database
queries no matter how many books there are.
"""
Book.objects.create(
author=Author.objects.create(name="test"),
title="test",
)
with self.assertNumQueries(1): # 👈🏻
self.client.get("/books/")If you try and run this test, you should see something like this:
======================================================================
FAIL: test_view_query_count (testapp.tests.BookListViewTestCase.test_view_query_count)
Make sure that listing books always triggers the same number of database
----------------------------------------------------------------------
Traceback (most recent call last):
File "./testapp/tests.py", line 17, in test_view_query_count
with self.assertNumQueries(1):
^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError: 2 != 1 : 2 queries executed, 1 expected
Captured queries were:
1. SELECT "testapp_book"."id", "testapp_book"."author_id", "testapp_book"."title" FROM "testapp_book"
2. SELECT "testapp_author"."id", "testapp_author"."name" FROM "testapp_author" WHERE "testapp_author"."id" = 1 LIMIT 21
----------------------------------------------------------------------
The test is failing as expected, showing that two database queries were executed when we expected only one. Great!
So let's fix the view:
def list_books(request):
data = {
"books": [
{
"id": book.pk,
"title": book.title,
"author": book.author.name,
}
for book in Book.objects.select_related("author") # 👈🏻
]
}
return JsonResponse(data)And now when we run the tests again, they pass! 🥳
Limitations of assertNumQueries
assertNumQueries gets you pretty far, but it can make the tests a little fragile.
Let's say for example we wanted to restrict our view to logged-in users. It's fairly easy to do this with
Django's @login_required decorator:
from django.contrib.auth.decorators import login_required
@login_required # 👈🏻
def list_books(request):
data = {
"books": [
{
"id": book.pk,
"title": book.title,
"author": book.author.name,
}
for book in Book.objects.select_related("author") # 👈🏻
]
}
return JsonResponse(data)Our test need to be updated as well to make sure we use a logged-in user:
def test_view_query_count(self):
"""
Make sure that listing books always triggers the same number of database
queries no matter how many books there are.
"""
Book.objects.create(
author=Author.objects.create(name="test"),
title="test",
)
self.client.force_login(User.objects.create_user("test")) # 👈🏻
with self.assertNumQueries(1):
self.client.get("/books/")Great, now let's run the tests again:
======================================================================
FAIL: test_view_query_count (testapp.tests.BookListViewTestCase.test_view_query_count)
Make sure that listing books always triggers the same number of database
----------------------------------------------------------------------
Traceback (most recent call last):
File "./testapp/tests.py", line 19, in test_view_query_count
with self.assertNumQueries(1):
^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError: 3 != 1 : 3 queries executed, 1 expected
Captured queries were:
1. SELECT "django_session"."session_key", "django_session"."session_data", "django_session"."expire_date" FROM "django_session" WHERE ("django_session"."expire_date" > '2025-08-07 13:51:14.193677' AND "django_session"."session_key" = '9fb11ct45wfepk5zwvty3a5pv5pgljsr') LIMIT 21
2. SELECT "auth_user"."id", "auth_user"."password", "auth_user"."last_login", "auth_user"."is_superuser", "auth_user"."username", "auth_user"."first_name", "auth_user"."last_name", "auth_user"."email", "auth_user"."is_staff", "auth_user"."is_active", "auth_user"."date_joined" FROM "auth_user" WHERE "auth_user"."id" = 1 LIMIT 21
3. SELECT "testapp_book"."id", "testapp_book"."author_id", "testapp_book"."title", "testapp_author"."id", "testapp_author"."name" FROM "testapp_book" INNER JOIN "testapp_author" ON ("testapp_book"."author_id" = "testapp_author"."id")
----------------------------------------------------------------------Oh no, it broke! I bet you didn't see that one coming. 😁
It turns out that in order to check whether the user is logged in before accessing the view, Django will execute some database queries:
- Look up the session information
- Look up the user information with the user id found in the session data
You might be tempted to just slap an assertNumQueries(3) in your test and call it a day. Sure,
that will work. But what happens the next time you make an unrelated change to the view? Maybe you change your
session backend to use cookies? Or maybe you retrieve a site-wide setting inside the view? ...
Having a brittle test like this which needs updating when changing seemingly unrelated things is a bummer, but what if I told you there was a better way?
An idea for a more robust approach
What if instead of asserting an exact number of database queries, we could do the following:
- Count how many queries are run with 1 instance in the db;
- Create a lot of new instances;
- Run the code again and make sure it triggers the same number of queries again.
Django's CaptureQueriesContext
Just a heads up: I'm about to suggest using an undocumented Django API. In general it's a pretty bad habit to
use a framework's undocumented features, but who doesn't like living dangerously? If that's not your jam 🍯,
you can either reimplement the equivalent functionality yourself, or manually count queries in
connection.queries.
The delicious forbidden fruit here is Django's juicy captureQueriesContext (located in
django.test.utils). It's another context manager like assertNumQueries but it works
a bit differently.
As the name hints at, this context manager will capture the queries executed inside its block and store them
on its captured_queries attribute (which will be a list of queries, as strings). It can be useful
on its own even outside of a test:
from django.db import connection
with captureQueriesContext(connection) as ctx:
list(MyModel.objects.all())
print(*ctx.captured_queries, sep="\n")
In our case, we don't care so much about the exact queries since we just want to count how many were executed.
Because captured_queries is a list, that means all we need is len(captured_queries):
from django.db import connection
with captureQueriesContext(connection) as ctx:
list(MyModel.objects.all())
query_count = len(ctx.captured_queries)We now have all the pieces we need to rewrite our test:
def test_view_query_count(self):
"""
Make sure that listing books always triggers the same number of database
queries no matter how many books there are.
"""
Book.objects.create(
author=Author.objects.create(name="test"),
title="test",
)
with CaptureQueriesContext(connection) as ctx:
self.client.get("/books/")
baseline = len(ctx.captured_queries)
for i in range(10):
# Here a tool like factory-boy or model-bakery would really help...
Book.objects.create(
author=Author.objects.create(name=f"test{i}"),
title="test{i}",
)
self.assertNumQueries(baseline):
self.client.get("/books/")And voilà 🎻, we now have a more robust test that won't break as often!