Django: is there a way to count SQL queries from an unit test?

DjangoDjango OrmDjango Testing

Django Problem Overview

I am trying to find out the number of queries executed by a utility function. I have written a unit test for this function and the function is working well. What I would like to do is track the number of SQL queries executed by the function so that I can see if there is any improvement after some refactoring.

def do_something_in_the_database():
    # Does something in the database
    # return result

class DoSomethingTests(django.test.TestCase):
    def test_function_returns_correct_values(self):
        self.assertEqual(n, <number of SQL queries executed>)

EDIT: I found out that there is a pending Django feature request for this. However the ticket is still open. In the meantime is there another way to go about this?

Django Solutions

Solution 1 - Django

Since Django 1.3 there is a assertNumQueries available exactly for this purpose.

One way to use it (as of Django 3.2) is as a context manager:

# measure queries of some_func and some_func2
with self.assertNumQueries(2):
    result = some_func()
    result2 = some_func2()

Solution 2 - Django

Vinay's response is correct, with one minor addition.

Django's unit test framework actually sets DEBUG to False when it runs, so no matter what you have in, you will not have anything populated in connection.queries in your unit test unless you re-enable debug mode. The Django docs explain the rationale for this as:

> Regardless of the value of the DEBUG setting in your configuration file, all Django tests run with DEBUG=False. This is to ensure that the observed output of your code matches what will be seen in a production setting.

If you're certain that enabling debug will not affect your tests (such as if you're specifically testing DB hits, as it sounds like you are), the solution is to temporarily re-enable debug in your unit test, then set it back afterward:

def test_myself(self):
    from django.conf import settings
    from django.db import connection
    settings.DEBUG = True
    connection.queries = []

    # Test code as normal

    settings.DEBUG = False

Solution 3 - Django

If you are using pytest, pytest-django has django_assert_num_queries fixture for this purpose:

def test_queries(django_assert_num_queries):
    with django_assert_num_queries(3):

Solution 4 - Django

If you don't want use TestCase (with assertNumQueries) or change settings to DEBUG=True, you can use context manager CaptureQueriesContext (same as assertNumQueries using).

from django.db import ConnectionHandler
from django.test.utils import CaptureQueriesContext

DB_NAME = "default"  # name of db configured in settings you want to use - "default" is standard
connection = ConnectionHandler()[DB_NAME]
with CaptureQueriesContext(connection) as context:
    ... # do your thing
num_queries = context.initial_queries - context.final_queries
assert num_queries == expected_num_queries

db settings

Solution 5 - Django

In modern Django (>=1.8) it's well documented (it's also documented for 1.7) here, you have the method reset_queries instead of assigning connection.queries=[] which indeed is raising an error, something like that works on django>=1.8:

class QueriesTests(django.test.TestCase):
    def test_queries(self):
        from django.conf import settings
        from django.db import connection, reset_queries

            settings.DEBUG = True
            # [... your ORM code ...]
            self.assertEquals(len(connection.queries), num_of_expected_queries)
            settings.DEBUG = False

You may also consider resetting queries on setUp/tearDown to ensure queries are reset for each test instead of doing it on finally clause, but this way is more explicit (although more verbose), or you can use reset_queries in the try clause as many times as you need to evaluate queries counting from 0.

Solution 6 - Django

Here is the working prototype of context manager withAssertNumQueriesLessThan

import json
from contextlib import contextmanager
from django.test.utils import CaptureQueriesContext
from django.db import connections

def withAssertNumQueriesLessThan(self, value, using='default', verbose=False):
    with CaptureQueriesContext(connections[using]) as context:
        yield   # your test will be run here
    if verbose:
        msg = "\r\n%s" % json.dumps(context.captured_queries, indent=4)
        msg = None
    self.assertLess(len(context.captured_queries), value, msg=msg)

It can be simply used in your unit tests for example for checking the number of queries per Django REST API call

    with self.withAssertNumQueriesLessThan(10):
        response = self.client.get('contacts/')
        self.assertEqual(response.status_code, 200)

Also you can provide exact DB using and verbose if you want to pretty-print list of actual queries to stdout

Solution 7 - Django

If you have DEBUG set to True in your (presumably so in your test environment) then you can count queries executed in your test as follows:

from django.db import connection

class DoSomethingTests(django.test.TestCase):
    def test_something_or_other(self):
        num_queries_old = len(connection.queries)
        num_queries_new = len(connection.queries)
        self.assertEqual(n, num_queries_new - num_queries_old)

Solution 8 - Django

If you want to use a decorator for that there is a nice gist:

import functools
import sys
import re
from django.conf import settings
from django.db import connection

def shrink_select(sql):
    return re.sub("^SELECT(.+)FROM", "SELECT .. FROM", sql)

def shrink_update(sql):
    return re.sub("SET(.+)WHERE", "SET .. WHERE", sql)

def shrink_insert(sql):
    return re.sub("\((.+)\)", "(..)", sql)

def shrink_sql(sql):
    return shrink_update(shrink_insert(shrink_select(sql)))

def _err_msg(num, expected_num, verbose, func=None):
    func_name = "%s:" % func.__name__ if func else ""
    msg = "%s Expected number of queries is %d, actual number is %d.\n" % (func_name, expected_num, num,)
    if verbose > 0:
        queries = [query['sql'] for query in connection.queries[-num:]]
        if verbose == 1:
            queries = [shrink_sql(sql) for sql in queries]
        msg += "== Queries == \n" +"\n".join(queries)
    return msg

def assertNumQueries(expected_num, verbose=1):

    class DecoratorOrContextManager(object):
        def __call__(self, func):  # decorator
            def inner(*args, **kwargs):
                handled = False
                    return func(*args, **kwargs)
                    handled = True
                    if not handled:
                        self.__exit__(None, None, None)
            return inner

        def __enter__(self):
            self.old_debug = settings.DEBUG
            self.old_query_count = len(connection.queries)
            settings.DEBUG = True

        def __exit__(self, type, value, traceback):
            if not type:
                num = len(connection.queries) - self.old_query_count
                assert expected_num == num, _err_msg(num, expected_num, verbose)
            settings.DEBUG = self.old_debug

    return DecoratorOrContextManager()


All content for this solution is sourced from the original question on Stackoverflow.

The content on this page is licensed under the Attribution-ShareAlike 4.0 International (CC BY-SA 4.0) license.

Content TypeOriginal AuthorOriginal Content on Stackoverflow
QuestionManoj GovindanView Question on Stackoverflow
Solution 1 - DjangoMitarView Answer on Stackoverflow
Solution 2 - DjangoJarret HardieView Answer on Stackoverflow
Solution 3 - DjangotvorogView Answer on Stackoverflow
Solution 4 - DjangoDaniel BartonView Answer on Stackoverflow
Solution 5 - DjangodaniusView Answer on Stackoverflow
Solution 6 - DjangopymenView Answer on Stackoverflow
Solution 7 - DjangoVinay SajipView Answer on Stackoverflow
Solution 8 - DjangoncopiyView Answer on Stackoverflow