如何Mock SQLAlchemy引擎，单元测试Python更新/插入查询函数

你遇到的返回Mock对象的问题，大概率是side_effect的设置方式不对——因为我们需要控制的是cursor的rowcount属性，而不是让execute方法返回Mock。下面具体说怎么实现：

先明确核心逻辑 #

你的update_worker_data函数大概是这个路子：

def update_worker_data(conn, worker_id, data):
    update_sql = "UPDATE workers SET name = %s WHERE id = %s"
    cursor = conn.cursor()
    cursor.execute(update_sql, (data["name"], worker_id))
    # 根据rowcount判断是否插入
    if cursor.rowcount == 0:
        insert_sql = "INSERT INTO workers (id, name) VALUES (%s, %s)"
        cursor.execute(insert_sql, (worker_id, data["name"]))
    conn.commit()

正确的测试写法 #

用pytest.mark.parameterize定义两个测试分支，再通过side_effect控制cursor的rowcount：

import pytest
from unittest.mock import Mock
from your_module import update_worker_data

@pytest.mark.parametrize("rowcounts, expected_executions", [
    ([1], 1),  # 更新分支：update影响1行，只执行1次execute
    ([0, 1], 2),  # 插入分支：update影响0行，触发插入，执行2次execute
])
def test_update_worker_data(rowcounts, expected_executions):
    # 初始化mock游标和连接
    mock_cursor = Mock()
    mock_conn = Mock()
    mock_conn.cursor.return_value = mock_cursor

    # 用迭代器控制每次execute后的rowcount
    rowcount_iter = iter(rowcounts)
    def execute_side_effect(*args):
        # 每次调用execute后，更新cursor的rowcount
        mock_cursor.rowcount = next(rowcount_iter)
    
    mock_cursor.execute.side_effect = execute_side_effect

    # 调用目标函数
    update_worker_data(mock_conn, 123, {"name": "Alice"})

    # 断言执行次数
    assert mock_cursor.execute.call_count == expected_executions

    # 分支断言：插入分支要验证调用了插入语句
    if expected_executions == 2:
        insert_call = mock_cursor.execute.call_args_list[1]
        assert "INSERT INTO workers" in insert_call[0][0]
    else:
        update_call = mock_cursor.execute.call_args_list[0]
        assert "UPDATE workers" in update_call[0][0]
    
    # 断言提交操作被触发
    assert mock_conn.commit.called

为什么之前会返回Mock对象？ #

如果你的side_effect直接设置成了Mock实例，那cursor.execute()会返回这个Mock，而不是去修改rowcount属性。我们需要的是通过side_effect的函数逻辑，动态修改mock游标rowcount的值——这才是触发插入分支的关键，而不是让execute返回某个对象。

总结 #

side_effect是处理这种多分支测试的标准方案，它能灵活控制mock对象的行为，完美适配你需要覆盖更新、插入两个分支的需求，只要把逻辑聚焦在修改rowcount属性上就行。

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