如何Mock SQLAlchemy的session.query并根据入参返回不同结果

你当前mock不生效的核心原因：
你mock的flask_sqlalchemy._QueryProperty.__get__是Model.query写法专属的底层实现，session.query是直接调用session实例的query方法，完全不走这个属性逻辑，所以肯定不会生效。

解决方案 #

1. 兼容mock session.query 写法 #

直接mock你业务代码实际使用的session对象的query方法即可，分两种情况：

• 如果业务代码里统一用的是全局db实例的session（比如from app.models import db然后用db.session.query），直接patch对应导入路径的db.session.query即可，比如patch路径写app.models.db.session.query
• 如果是测试用fixture生成的session实例，直接给这个实例的query属性赋值为MagicMock更简单，不需要走patch逻辑。

2. 按查询入参返回不同mock值 #

利用MagicMock的side_effect特性，写自定义判断逻辑的函数，根据query接收的入参判断要查询的模型，返回对应结果即可。

完整测试代码示例：

from unittest import mock
from app.models import User, Book, check_user_own_book

def test_mock_session_query(fixture_user, fixture_book, session):
    # 自定义query的side_effect，根据入参返回对应mock结果
    def mock_query_side_effect(*args):
        mock_query = mock.MagicMock()
        # 识别查询的是User模型
        if any(hasattr(arg, 'class_') and arg.class_ == User for arg in args):
            mock_query.filter.return_value.one_or_none.return_value = fixture_user
        # 识别查询的是Book模型
        elif any(hasattr(arg, 'class_') and arg.class_ == Book for arg in args):
            mock_query.filter.return_value.one_or_none.return_value = fixture_book
        # 其他查询默认返回None
        else:
            mock_query.filter.return_value.one_or_none.return_value = None
        return mock_query

    # 给fixture的session绑定mock后的query方法
    session.query = mock.MagicMock(side_effect=mock_query_side_effect)

    # 兼容Model.query写法的mock，调整原来的patch逻辑支持按模型返回不同值
    with mock.patch('flask_sqlalchemy._QueryProperty.__get__') as mock_model_get:
        def model_query_side_effect(instance, owner):
            mock_q = mock.MagicMock()
            if owner == User:
                mock_q.filter.return_value.one_or_none.return_value = fixture_user
            elif owner == Book:
                mock_q.filter.return_value.one_or_none.return_value = fixture_book
            return mock_q
        mock_model_get.side_effect = model_query_side_effect

        # 两种查询写法都会返回对应结果
        query_user_with_session = session.query(User.id, User.username).filter(User.id==1000).one_or_none()
        assert query_user_with_session == fixture_user

        query_book_with_session = session.query(Book.id, Book.name).filter(Book.id==10000).one_or_none()
        assert query_book_with_session == fixture_book

        query_user_with_class = User.query.filter(User.id==1).one_or_none()
        assert query_user_with_class == fixture_user

        query_book_with_class = Book.query.filter(Book.id==1).one_or_none()
        assert query_book_with_class == fixture_book

        # 业务函数调用不会出现返回值串的问题
        has_own_book = check_user_own_book(user_id=1)

补充说明 #

如果需要更细粒度的控制，比如还要根据filter的条件、是否调用first()/one()/all()等方法返回不同结果，只需要在side_effect的逻辑里再加分层判断即可。
另外你贴的示例代码里有两处笔误：模型定义里owned_by = =多了一个等号，查询Book的filter条件里owned_by == user.id应该写为Book.owned_by == user.id，不然实际运行会报错。

内容的提问来源于stack exchange，提问作者Bao Tran