Search code examples
pythondjangopytestpytest-django

Get overview of SQL query count for every test in a test suite


I have a large Django application with tons of tests that need SQL query optimizations.

I'm using pytest-django to run my tests.

I don't want to add assertNumQueries or django-assert-num-queries for each test individually, but instead produce an overview about how many SQL queries each one of all of the tests are triggering, to know which code needs the most optimization, like this:

test                         | number of queries
------------------------------------------------
test_a.py::test_my_function1 |  5
test_a.py::test_my_function3 |  2
test_a.py::test_my_function5 |  7

Is it possible to configure a pytest hook in conftest.py which counts the number of SQL queries for each (DB) test and shows them in result - without the need to modify the source of my tests (like adding decorators)?

My naive approach would be to use these hooks, and somehow get access to the database connection before and after each test:

def pytest_runtest_call(item):
    pass

def pytest_runtest_teardown(item, nextitem):
    return True


Solution

  • For recording the queries count, an autouse fixture will suffice. In the below example, I store the count in a queries_count dict under the config object:

    @pytest.fixture(autouse=True)
    def record_query_count(request):
        from django.test.utils import CaptureQueriesContext
        from django.db import connection
    
        with CaptureQueriesContext(connection) as context:
            yield
        num_performed = len(context)
    
        if not hasattr(request.config, "query_counts"):
            request.config.query_counts = dict()
        request.config.query_counts[request.node.nodeid] = num_performed
    

    To output the results, I add a custom section in a custom impl of the pytest_terminal_summary hook. Put the following code in a file named conftest.py in the project or tests root dir:

    import os
    
    def pytest_terminal_summary(terminalreporter, exitstatus, config):
        content = os.linesep.join(
            f'{nodeid:40} | {num_queries}'
            for nodeid, num_queries in config.query_counts.items()
        )
        terminalreporter.ensure_newline()
        terminalreporter.section('number of queries', sep='-', bold=True)
        terminalreporter.line(content)
    

    Running the tests will now yield:

    tests/test_db.py ...                                                                [100%]
    
    ----------------------------------- number of queries ------------------------------------
    tests/test_db.py::test_add_entry      | 2
    tests/test_db.py::test_remove_entry   | 2
    tests/test_db.py::test_no_entries     | 1
    =================================== 3 passed in 0.26s ====================================