How to count sqlalchemy queries in unit tests

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

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.

Slightly modified version of @omar-tarabai's solution that removes the event listener when exiting the context:

from sqlalchemy import event

class QueryCounter(object):
    """Context manager to count SQLALchemy queries."""

    def __init__(self, connection):
        self.connection = connection.engine
        self.count = 0

    def __enter__(self):
        event.listen(self.connection, "before_cursor_execute", self.callback)
        return self

    def __exit__(self, *args, **kwargs):
        event.remove(self.connection, "before_cursor_execute", self.callback)

    def callback(self, *args, **kwargs):
        self.count += 1

Usage:

with QueryCounter(session.connection()) as counter:
     session.query(XXX).all()
     session.query(YYY).all()

print(counter.count)  # 2

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.