CODEWITHSUNDEEP
pythondjangopytestpytest-django
Wolkenarchitekt
Wolkenarchitekt

Reputation: 21248

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

Upvotes: 6

Views: 1638

Answers (1)

hoefling
hoefling

Reputation: 66501

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 ====================================

Upvotes: 8

Related Questions