如何用MagicMock模拟SQLAlchemy主查询中的子查询对象？

问题场景 #

我们通过如下方式动态构建SQLAlchemy查询：

def get_foo_results():
    included_bars_query = _get_included_bars_query(70)
    foo_results_query = (
        session.query(
            Foo.baz.label("baz"),
            Foo.qux.label("qux"),
            func.sum(Foo.dollars).label("dollars"),
        )
        .filter(
            Foo.bar.in_(included_bars_query),
        )
        .group_by(Foo.baz, Foo.qux)
    )
    return foo_results_query

def _get_included_bars_query(batch_number):
    valid_bars_query = session.query(Bars.bar_id.label("BAR_ID"))
    valid_bars_query = valid_bars_query.filter(Bars.batch_number > batch_number)
    return valid_bars_query

我们希望通过Mock进行规范化的针对性单元测试，避免为每个测试搭建完整数据库。已创建MagicMock对象来模拟get_foo_results中的_get_included_bars_query子查询，但不知如何让该Mock在主查询中生效，尝试了以下几种写法：

bars_query.return_value.subquery.return_value = [1, 2,]

bars_query.return_value.all.return_value = [1, 2,]

bars_query.return_value.subquery.return_value = [(1,), (2,),]

感觉思路接近，但因对SQLAlchemy理解不足未能成功，恳请解释问题出在哪里？

问题根源与解决方法 #

核心问题 #

当你在SQLAlchemy中使用Foo.bar.in_(subquery)时，SQLAlchemy并不会直接执行子查询并取结果列表，而是会把这个子查询作为嵌套SQL语句嵌入主查询中。所以你Mock的all()或者subquery()返回值根本不会被主查询调用——主查询需要的是一个能被编译成SQL的查询对象，而不是直接的结果列表。

正确Mock方式 #

你需要让_get_included_bars_query的Mock返回一个模拟的查询对象，这个对象要能被SQLAlchemy的in_()表达式正确处理。具体实操方式：

from unittest.mock import MagicMock, patch

def test_get_foo_results():
    # 模拟子查询对应的bar_id列表
    mock_bar_ids = [1, 2]
    # 创建模拟的子查询对象
    mock_subquery = MagicMock()
    # 关键：让子查询对象具备SQLAlchemy Query所需的列属性
    mock_subquery.columns.BAR_ID = MagicMock()
    
    # Mock _get_included_bars_query，让它返回模拟的子查询对象
    with patch('__main__._get_included_bars_query', return_value=mock_subquery):
        # 同时Mock主查询的session和返回结果
        mock_session = MagicMock()
        mock_query = MagicMock()
        # 设置主查询执行后的预期返回值
        mock_query.filter.return_value.group_by.return_value.all.return_value = [
            {"baz": "test1", "qux": "test2", "dollars": 100}
        ]
        mock_session.query.return_value = mock_query
        
        # 确保get_foo_results使用的是mock_session（可通过依赖注入或全局Mock实现）
        results = get_foo_results().all()
        # 验证结果符合预期
        assert results == [{"baz": "test1", "qux": "test2", "dollars": 100}]
        # 验证子查询是否被正确传入in_()条件
        mock_query.filter.assert_called_once()
        call_condition = mock_query.filter.call_args[0][0]
        # 可进一步验证call_condition是否为Foo.bar.in_(mock_subquery)

之前写法无效的原因 #

• bars_query.return_value.subquery.return_value = [1,2]：SQLAlchemy不会主动调用subquery()方法，_get_included_bars_query返回的已经是Query对象，主查询直接将其作为子查询嵌入，不会触发subquery()调用。
• bars_query.return_value.all.return_value = [1,2]：主查询不会执行子查询的all()方法，它只会把子查询转化为嵌套SQL，这个返回值永远不会被用到。
• bars_query.return_value.subquery.return_value = [(1,), (2,)]：同样不会触发subquery()调用，且SQLAlchemy需要的是可编译为SQL的查询结构，不是原始数据列表。

内容的提问来源于stack exchange，提问作者Nathaniel Rink