Python单元测试@patch下assert_called_with报AssertionError排查

问题现象 #

使用@patch装饰器完成MySQL连接相关mock打桩时，patch逻辑看似生效，预期被打桩的execute方法会被分4次传入不同参数调用，但实际运行抛出断言错误，无法定位根因。

相关代码 #

测试代码 #

class testMysqlBase(mySQLTestCase):
    def setUp(self):
        super().setUp()

    def tearDown(self) -> None:
        return super().tearDown()

    @patch("mysql.connector.connection.MySQLConnection")
    def test_exception_create_db_called_if_no_db(self, mysql_cnx) -> None:
        mysql_db_ex = MysqlDBExecutor(mysql_cnx, "test_db")
        mysql_cnx.cursor.return_value.__enter__.return_value = mysql_cnx.cursor

        mysql_cnx.cursor.return_value.__enter__.return_value.execute.side_effect = [
            (mysql.connector.errors.DatabaseError(errno=errorcode.ER_BAD_DB_ERROR)),
            None,
            None,
            None,
        
            mysql_db_ex.execute("test")
            mysql_cnx.cursor.execute.assert_called_with(
               ["USE test_db;",
               "CREATE DATABASE IF NOT EXISTS test_db",
              "USE test_db;",
               "test"]
        )

被测业务代码 #

from re import I

import mysql.connector
from mysql.connector import errorcode


class MysqlDBExecutor:

    logger: Logger = getLogger(__name__)

    def __init__(self, cnx: mysql.connector.connection, db: str):
        self.cnx = cnx
        self.db = db

    def execute(self, cmd) -> None:

        with self.cnx.cursor() as mycursor:
            try:
                mycursor.execute("USE " + self.db + ";")
            except mysql.connector.Error as err:
                if err.errno == errorcode.ER_BAD_DB_ERROR:
                    self.create_mysql_db()
                    mycursor.execute("USE " + self.db + ";")
                else:
                    raise err(
                        f"Exception thrown in mysqlDBExecutor.execute on {self.db}"
                    )

            try:           
                mycursor.execute(cmd)
            except mysql.connector.Error as err:
                raise err(f"Exception thrown in mysqlDBExecutor.execute on {self.db}")

    def create_mysql_db(self) -> None:

        mysql_query = "CREATE DATABASE IF NOT EXISTS {}".format(self.db)

        with self.cnx.cursor() as mycursor:
            try:
                mycursor.execute(mysql_query)
            except mysql.connector.Error as err:
                raise err(
                    f"mysql Exception thrown in mysqlExecutor creating db{self.db}"
                )
            except Exception:
                raise

报错信息 #

E           AssertionError: expected call not found.
E           Expected: execute(['USE test_db;', 'CREATE DATABASE IF NOT EXISTS test_db', 'USE test_db;', 'test'])
E           Actual: execute('test')

预期逻辑 #

访问不存在的数据库、调用mysql_db_ex.execute方法时，被mock的mysql_cnx游标首次调用execute会触发errno为ER_BAD_DB_ERROR的DatabaseError；异常被捕获后会调用self.create_mysql_db()方法，内部调用游标执行建库SQL；之后重新执行USE语句切换到目标库，最后执行传入的cmd命令。整个流程execute方法应当被调用4次。

根因分析 #

代码里存在3个核心错误，直接导致运行结果不符合预期：

• 括号闭合错误，执行顺序完全错乱
  side_effect赋值的列表字面量没有正确闭合，mysql_db_ex.execute("test")和后续断言代码都被写在了side_effect的列表内部。Python赋值语句会先计算右侧表达式的值再完成赋值，也就是说：
  • 执行测试方法时，会先求值side_effect =右侧的列表，此时会直接触发mysql_db_ex.execute("test")的调用
  • 调用execute时，side_effect还没完成赋值，首次执行USE test_db语句不会抛出预期的数据库不存在异常，自然不会走到建库分支
  • 整个流程只会执行2次execute调用：第一次USE语句、第二次传入的test语句，断言时只会捕获到最后一次execute('test')的调用记录
• Mock对象绑定逻辑错误
  代码将上下文管理器的返回值__enter__.return_value绑定为mysql_cnx.cursor（即mock的cursor方法本身），而create_mysql_db方法内部会重新调用self.cnx.cursor()获取新游标，两次上下文管理器拿到的不是同一个mock实例，调用记录会被分散，无法统一统计。
• 断言方法使用错误
  assert_called_with仅会校验mock对象最后一次被调用的参数，且直接传入了一个4元素列表作为参数，等价于断言execute被调用时入参是这个列表，和实际每次传入单个字符串的调用逻辑完全不符。

修复方案 #

按以下步骤修改测试代码即可：

• 修正括号位置，闭合side_effect的列表，将业务调用、断言逻辑移到side_effect赋值完成之后
• 统一游标mock实例，确保所有cursor()上下文管理器调用都返回同一个mock游标，避免调用记录分散
• 替换断言方法为assert_has_calls，配合call对象按顺序校验所有调用的入参

修复后的测试代码参考：

from unittest.mock import patch, call

# 其他导入省略...

class testMysqlBase(mySQLTestCase):
    def setUp(self):
        super().setUp()

    def tearDown(self) -> None:
        return super().tearDown()

    @patch("mysql.connector.connection.MySQLConnection")
    def test_exception_create_db_called_if_no_db(self, mysql_cnx) -> None:
        mysql_db_ex = MysqlDBExecutor(mysql_cnx, "test_db")
        # 统一绑定游标mock实例，所有cursor()调用返回同一个mock
        mock_cursor = mysql_cnx.cursor.return_value.__enter__.return_value
        # 正确配置side_effect，闭合列表
        mock_cursor.execute.side_effect = [
            mysql.connector.errors.DatabaseError(errno=errorcode.ER_BAD_DB_ERROR),
            None,
            None,
            None,
        ]
        # side_effect配置完成后再执行业务方法
        mysql_db_ex.execute("test")
        # 用assert_has_calls按顺序校验所有调用
        mock_cursor.execute.assert_has_calls([
            call("USE test_db;"),
            call("CREATE DATABASE IF NOT EXISTS test_db"),
            call("USE test_db;"),
            call("test")
        ])

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