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

Question

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

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

当我运行以上时，我得到以下

失败 tests/unit/test_code.py::test_does_table_exist - 断言 None == True

嘲笑有时让我很困惑。我需要确保 cur.execute 被模拟为 return 1 回来。但是 None 被函数 return 编辑了。有人能告诉我如何正确地模拟这个吗？

Answer 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

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

Mocking pymysql.connect calls in Python unit test

python

database

unit-testing

mocking

pymysql