在 contextManager 下使用 unittest 模拟 MysqlConnection

Question

我正在尝试测试一种在 with 块内检索 mysql 连接的方法，但是当我尝试从中获取模拟游标时，并且从游标执行的结果未正确模拟。

这里是获取连接的方法

def has_pending_task(self, host, port):
    with self.get_mysql_connection(host, port) as conn:
        cursor = conn.cursor()
        table = self.config.get_table()
        cursor.execute(
            'select count(*) from {} where status in ("PENDING", "RUNNING") and host=%s'.format(
                table), host)
        count = cursor.fetchone()
        cursor.close()

        return count > 0

@contextmanager
def get_mysql_connection(self, host, port):
    db_credentials = self.config.get_credentials()
    schema = self.config.get_schema()
    if not db_credentials:
        logging.error(
            "There was an error getting db info for cluster {}. Please review the config".format(self.cluster))
        return None

    conn = None
    try:
        conn = MySQLConnection(user=db_credentials[0], password=db_credentials[1],
                                          host=host, port=port, database=schema)
    except Error as error:
        logging.error("Couldn't create a mysql connection: {}".format(error.msg))

    return conn

测试看起来像这样：

def test_has_no_pending_tasks(self):
    with mock.patch("my_package.my_class.Config") as config_mock:
        config = config_mock.return_value
        config_mock.get_instance.return_value = config
        config.get_table.return_value = "table"

        with mock.patch("mysql.connector.connection.MySQLConnection") as mysql_mock:
            mysql = mysql_mock.return_value
            with mock.patch('package.MyClass.get_mysql_connection') as conn_mock:
                conn_mock.return_value = mysql
                cursor = mysql_mock.cursor.return_value
                cursor.fetchone.return_value = 0

                Orchestrator("host", 3306).has_pending_tasks("host", 3306)

                conn_mock.assert_called_once_with("host", 3306)

问题在于，由于它没有正确地模拟事物，测试函数何时返回，count 是 MagicMock 的一个实例并导致测试崩溃。

知道我缺少什么部分吗？

谢谢！

Answer 1

根据Python: Mocking a context manager，问题出在上下文管理器上。像

这样模拟 get_mysql_connection() 方法

onn_mock.return_value.__enter__.return_value = mysql

它完美地做到了