Django：有没有办法计算单元测试中的 SQL 查询数量？

Question

我试图找出实用函数执行的查询数量。 我已经为此函数编写了单元测试，并且该函数运行良好。 我想要做的是跟踪该函数执行的 SQL 查询的数量，以便我可以查看一些重构后是否有任何改进。

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

编辑：我发现有一个待处理的 Django 功能请求。 然而，门票仍然开放。 与此同时，还有其他方法可以解决这个问题吗？

Answer 1

如果你想使用装饰器，有一个很好的要点：

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
            @functools.wraps(func)
            def inner(*args, **kwargs):
                handled = False
                try:
                    self.__enter__()
                    return func(*args, **kwargs)
                except:
                    self.__exit__(*sys.exc_info())
                    handled = True
                    raise
                finally:
                    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()

Answer 2

如果您在 settings.py 中将 DEBUG 设置为 True（大概在您的测试环境中如此），那么您可以对测试中执行的查询进行计数，如下所示：

from django.db import connection

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

Answer 3

这是上下文管理器 withAssertNumQueriesLessThan 的工作原型，

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

@contextmanager
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)
    else:
        msg = None
    self.assertLess(len(context.captured_queries), value, msg=msg)

它可以简单地在您的单元测试中使用，例如检查每个 Django REST API 调用的查询数量 数据库

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

如果您想要将实际查询的列表漂亮地打印到标准输出，您还可以using和verbose提供精确的

Answer 4

在现代 Django (>=1.8) 中，它有详细的文档记录（1.7 也有文档记录）此处，您可以使用方法reset_queries，而不是分配connection.queries=[ ] 这确实引发了一个错误，类似于 django>=1.8 上的情况：

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

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

您还可以考虑在 setUp/tearDown 上重置查询，以确保为每个测试重置查询，而不是在 finally 子句上执行此操作，但这way 更明确（尽管更详细），或者您可以根据需要在 try 子句中使用 reset_queries 多次来评估从 0 开始计数的查询。

Answer 5

如果您不想使用 TestCase （使用 assertNumQueries）或将设置更改为 DEBUG=True，您可以使用上下文管理器 CaptureQueriesContext （与 assertNumQueries 使用）。

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

数据库设置

Answer 6

如果您使用 pytest，则 pytest-django 具有 django_assert_num_queries 用于此目的的固定装置：

def test_queries(django_assert_num_queries):
    with django_assert_num_queries(3):
        Item.objects.create('foo')
        Item.objects.create('bar')
        Item.objects.create('baz')

Answer 7

Vinay 的回答是正确的，但有一点补充。

Django 的单元测试框架实际上在运行时将 DEBUG 设置为 False，因此无论您在 settings.py 中有什么，您都不会在 connection.queries 中填充任何内容。单元测试，除非您重新启用调试模式。 Django 文档将基本原理解释为：

无论配置文件中 DEBUG 设置的值如何，所有 Django 测试都以 DEBUG=False 运行。 这是为了确保观察到的代码输出与生产环境中看到的输出相匹配。

如果您确定启用调试不会影响您的测试（例如，如果您专门测试数据库命中，就像听起来那样），解决方案是在单元测试中暂时重新启用调试，然后设置它之后回来：

def test_myself(self):
    from django.conf import settings
    from django.db import connection

    settings.DEBUG = True
    connection.queries = []

    # Test code as normal
    self.assert_(connection.queries)

    settings.DEBUG = False

Answer 8

从 Django 1.3 开始，有一个 assertNumQueries 正是用于此目的。

使用它的一种方法（从 Django 3.2 开始）是作为上下文管理器：

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