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 功能请求 。但是该请求仍然保持开放状态。与此同时，有其他方法可以解决这个问题吗？

- Manoj Govindan

Answer 1

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

- Mitar

---

https://docs.djangoproject.com/en/1.4/topics/testing/#django.test.TestCase.assertNumQueries - BenH

如果您需要使用“AssertNumQueriesLess”，欢迎访问https://dev59.com/jHM_5IYBdhLWcg3wt1vD#59089020。 - pymen

Answer 2

Vinay的回答是正确的，只有一个小补充。

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

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

如果您确定启用调试不会影响您的测试（例如，如果您正在专门测试DB访问情况，就像您所描述的那样），则解决方案是在单元测试中暂时重新启用调试，然后在之后将其设置回来：

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

- Jarret Hardie

---

谢谢。切换DEBUG就是我需要做的。 :) - Manoj Govindan

另外需要说明的是，你应该将测试代码放在try:块中，并将设置settings.DEBUG = False放在对应的finally:块中。这样，如果此测试失败，您的其他测试不会因DEBUG设置而受到“污染”。 - SmileyChris

你可以使用 connection.use_debug_cursor = True 来替代 settings.DEBUG = True。在我看来，这将是一个更本地化的解决方案。 - Oduvan

如果您想要针对每个测试记录查询日志，那么在setUp()和tearDown()中设置DEBUG=True是理想的。如果您希望在某些测试中以类似于生产数据库的方式运行，并仅记录特定测试中的查询，则最好在测试函数本身中更改设置。我想这完全取决于您的测试目标。 - Jarret Hardie

这是一种非常糟糕的更改设置的方式（正如@SmileyChris所暗示的），Django有很多临时更改设置而不会污染其他测试的方法（至少自Django 1.4以来，@override_settings装饰器已经存在）。 - Izkata

Answer 3

如果您正在使用 pytest，pytest-django 为此目的提供了 django_assert_num_queries fixture：

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

- tvorog

Answer 4

如果您不想使用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

数据库设置

- Daniel Barton

---

CaptureQueriesContext是一个被严重低估的测试上下文处理程序。您可以深入了解ORM执行的所有内容及其原因。 - GDorn

Answer 5

在现代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子句中执行此操作，但这种方式更加明确（尽管更冗长），或者你可以在try子句中使用reset_queries，根据需要多次重置查询计数从0开始。

- danius

Answer 6

这是 context manager 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)

如果您希望在stdout上漂亮打印实际查询的列表，您可以提供确切的DB using和verbose

- pymen

Answer 7

如果在你的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)

- Vinay Sajip

---

谢谢。我试过了，但奇怪的是，在调用函数之前和之后，len(connection.queries)都为零！我尝试将函数调用替换为直接调用MyModel.objects.filter()，但仍然没有运气。顺便说一下，我正在使用Django 1.1。 - Manoj Govindan

更新：如果我使用iPython交互式执行函数，该机制可以正常工作。当然，这是针对开发数据库而不是短暂测试数据库的。这种差异是否与Django在事务中执行测试的方式有关？ - Manoj Govindan

在Django测试中，DEBUG默认设置为False。这是因为您想要测试您的实际环境。 - DylanYoung

使用以下代码可以在Django测试中覆盖设置并获取查询次数：with self.settings(DEBUG=True): ... num_queries = len(connection.queries)参见文档 - djvg

Answer 8

如果你想为此使用装饰器，可以参考这个好的代码片段：

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

- ncopiy