在 Python 单元测试中模拟 pymysql.connect 调用答案

【问题标题】：Mocking pymysql.connect calls in Python unit test在 Python 单元测试中模拟 pymysql.connect 调用
【发布时间】：2021-12-02 13:51:24
【问题描述】：

我正在尝试模拟一个看起来像这样的函数。

code.app.py

import pymysql

# conn returns a pymysql.connect connection
conn = get_conn(db_port, db_username, db_name)


def does_table_exist(conn, table_name):
    
    sql = f"SELECT * FROM information_schema.tables WHERE table_name='{table_name}';"

    try:
        table_exists = None
        with conn.cursor() as cur:
            sql_result = cur.execute(sql)
            
            if sql_result == 1:
                table_exists = True
            if sql_result == 0:
                table_exists = False

            return table_exists
    except Exception as error:
        logger.error(error)
        raise
    finally:
        cur.close()

我的测试文件如下所示。

test_code.py

def test_check_table_exists(init):
    with patch('code.app.pymysql.connect') as mock_connect:
        cursor = MagicMock()
        execute = Mock(return_value=1)
        cursor.return_value = execute
        mock_connect.return_value.__enter__.return_value = cursor
        
        # Here init is a pytest fixture that imports the module
        response = init.does_table_exist(mock_connect, 'test')
        assert response == True

当我运行上述内容时，我得到以下内容

失败的测试/单元/test_code.py::test_does_table_exist - 断言无 == True

嘲笑有时完全让我感到困惑。我需要确保 cur.execute 被模拟以返回 1 。但是函数没有返回 None 。有人可以指点我如何正确地模拟这个。

【问题讨论】：

标签： python database unit-testing mocking pymysql

【解决方案1】：

我能够通过以下方式使其工作。

@pytest.mark.parametrize('data, output', [
    (1, True), (0, False)
])
def test_does_table_exist(init, data, output):
    with patch('code.app.pymysql.connect') as mock_connect:
        cursor = MagicMock()
        cursor.execute.return_value = data
        mock_connect.return_value.cursor.return_value.__enter__.return_value = cursor
        t = pymysql.connect()
        response = init.does_table_exist(t, 'test')
        assert response == output

【讨论】：