在MysqlConnection下使用unittest模拟contextManager

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

它做得很好