如何在单元测试中统计sqlalchemy查询

Question

在Django中,我经常声明应该进行的查询数量,以便单元测试能够捕获新的N + 1查询问题

from django import db
from django.conf import settings
settings.DEBUG=True

class SendData(TestCase):
    def test_send(self):
        db.connection.queries = []
        event = Events.objects.all()[1:]
        s = str(event) # QuerySet is lazy, force retrieval
        self.assertEquals(len(db.connection.queries), 2)

在SQLAlchemy中,通过echo在引擎上设置标志来启用对STDOUT的跟踪

engine.echo=True

编写计算SQLAlchemy查询数量的测试的最佳方法是什么？

class SendData(TestCase):
    def test_send(self):
        event = session.query(Events).first()
        s = str(event)
        self.assertEquals( ... , 2)

Answer 1

使用SQLAlchemy 核心事件来记录/跟踪执行的查询(您可以从单元测试中附加它,这样它们就不会影响您在实际应用程序上的性能:

event.listen(engine, "before_cursor_execute", catch_queries)

现在你编写函数catch_queries,其中的方式取决于你的测试方式.例如,您可以在test语句中定义此函数:

def test_something(self):
    stmts = []
    def catch_queries(conn, cursor, statement, ...):
        stmts.append(statement)
    # Now attach it as a listener and work with the collected events after running your test

上述方法只是一个灵感.对于扩展的情况,您可能希望拥有在每次测试后清空的事件的全局缓存.原因是在0.9(当前开发)之前没有用于删除事件侦听器的API.因此,创建一个访问全局列表的全局侦听器.

Answer 2

我为此创建了一个上下文管理器类:

class DBStatementCounter(object):
    """
    Use as a context manager to count the number of execute()'s performed
    against the given sqlalchemy connection.

    Usage:
        with DBStatementCounter(conn) as ctr:
            conn.execute("SELECT 1")
            conn.execute("SELECT 1")
        assert ctr.get_count() == 2
    """
    def __init__(self, conn):
        self.conn = conn
        self.count = 0
        # Will have to rely on this since sqlalchemy 0.8 does not support
        # removing event listeners
        self.do_count = False
        sqlalchemy.event.listen(conn, 'after_execute', self.callback)

    def __enter__(self):
        self.do_count = True
        return self

    def __exit__(self, *_):
        self.do_count = False

    def get_count(self):
        return self.count

    def callback(self, *_):
        if self.do_count:
            self.count += 1