How to count sqlalchemy queries in unit tests

Question

In Django I often assert the number of queries that should be made so that unit tests catch new N+1 query problems

from django import db
from django.conf import          


        

Answer 1

what about the approach of using flask_sqlalchemy.get_debug_queries() btw. this is the methodology used by internal of Flask Debug Toolbar check its source

from flask_sqlalchemy import get_debug_queries


def test_list_with_assuring_queries_count(app, client):
    with app.app_context():
        # here generating some test data
        for _ in range(10):
            notebook = create_test_scheduled_notebook_based_on_notebook_file(
                db.session, owner='testing_user',
                schedule={"kind": SCHEDULE_FREQUENCY_DAILY}
            )
            for _ in range(100):
                create_test_scheduled_notebook_run(db.session, notebook_id=notebook.id)
    with app.app_context():
        # after resetting the context call actual view we want asserNumOfQueries
        client.get(url_for('notebooks.personal_notebooks'))
        assert len(get_debug_queries()) == 3

keep in mind that for having reset context and count you have to call with app.app_context() before the exact stuff you want to measure.

Answer 2

Use SQLAlchemy Core Events to log/track queries executed (you can attach it from your unit tests so they don't impact your performance on the actual application:

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

Now you write the function catch_queries, where the way depends on how you test. For example, you could define this function in your test statement:

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

The above method is just an inspiration. For extended cases you'd probably like to have a global cache of events that you empty after each test. The reason is that prior to 0.9 (current dev) there is no API to remove event listeners. Thus make one global listener that accesses a global list.

Answer 3

I've created a context manager class for this purpose:

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