Python测试：如何Mock SQLAlchemy的查询操作？

问题背景 #

已完成数据库连接和验证逻辑的Mock，但不清楚如何Mock SQL插入操作。以下是相关代码，run_insert函数会在user_register中被调用，需要正确的Mock方法。

query_functions.py #

CONST_INSERT = "INSERT INTO user (user_name, password, email_address, dob, address, uid ) VALUES(:user_name , :password, :email_address, :dob, :address, :uid)"

def connection():
    engine = create_engine('mysql://root:admin@172.17.0.2:3306/flask', poolclass = NullPool)
    conn = engine.connect()
    return conn

def run_insert(data):
    conn = connection()
    sql = text(CONST_INSERT)
    conn.execute(sql, data)
    conn.close()
    return ("executed")

test.py（已尝试的代码） #

@patch.object(query_functions, 'connection')   
def test_service_reg_return_correct_data(self, mock_sqlalchmey):
    data = {                                                               
        "user_name": "testuser56",
        "password": "123456",
        "email_address": "testuser56@example.com",
        "dob": "2022-06-07 00:00:00",
        "address": "Ahmedabad"
    }
    
    m1 = mock.MagicMock()
    m1.get_json.return_value = data

    with mock.patch("backened.routes.register.request", m1):
        response = user_register() 
        self.assertTrue(mock_sqlalchmey.called)
        self.assertEqual(response.json, {'Message': 'New user Created'},201) 

user_register函数调用逻辑 #

def user_register():
     # 部分代码省略
    validate = validate()
     # 部分代码省略
    insert = run_insert(data)  # 注：原代码漏传data参数，修正为实际调用逻辑

两种可行的Mock方案 #

方案1：直接Mock run_insert函数 #

这是最直接的方式，只需验证user_register是否正确触发插入逻辑，无需关心内部数据库连接细节。

修改测试代码如下：

@patch.object(query_functions, 'run_insert')   
def test_service_reg_return_correct_data(self, mock_run_insert):
    # 设置Mock返回值
    mock_run_insert.return_value = "executed"
    
    data = {                                                               
        "user_name": "testuser56",
        "password": "123456",
        "email_address": "testuser56@example.com",
        "dob": "2022-06-07 00:00:00",
        "address": "Ahmedabad"
    }
    
    m1 = mock.MagicMock()
    m1.get_json.return_value = data

    with mock.patch("backened.routes.register.request", m1):
        response = user_register() 
        # 验证run_insert是否被调用，且传入正确参数
        mock_run_insert.assert_called_once_with(data)
        self.assertEqual(response.json, {'Message': 'New user Created'})
        self.assertEqual(response.status_code, 201)

方案2：Mock数据库连接对象的execute方法 #

如果需要验证插入语句的执行细节（比如SQL是否正确、参数是否匹配），可以完善对connection的Mock：

@patch.object(query_functions, 'connection')   
def test_service_reg_return_correct_data(self, mock_connection):
    # 创建连接对象Mock并设置返回值
    mock_conn = mock.MagicMock()
    mock_connection.return_value = mock_conn
    
    data = {                                                               
        "user_name": "testuser56",
        "password": "123456",
        "email_address": "testuser56@example.com",
        "dob": "2022-06-07 00:00:00",
        "address": "Ahmedabad"
    }
    
    m1 = mock.MagicMock()
    m1.get_json.return_value = data

    with mock.patch("backened.routes.register.request", m1):
        response = user_register() 
        # 验证连接方法被调用
        mock_connection.assert_called_once()
        # 验证execute方法被调用，且传入正确的SQL和参数
        mock_conn.execute.assert_called_once()
        called_sql, called_data = mock_conn.execute.call_args[0]
        self.assertEqual(str(called_sql), query_functions.CONST_INSERT)
        self.assertEqual(called_data["user_name"], data["user_name"])
        self.assertEqual(called_data["password"], data["password"])
        # 验证连接被关闭
        mock_conn.close.assert_called_once()
        
        self.assertEqual(response.json, {'Message': 'New user Created'})
        self.assertEqual(response.status_code, 201)

注意事项 #

• 原user_register代码中run_insert()漏传data参数，实际调用需修正为run_insert(data)，否则测试会报错。
• 如果uid是在user_register中生成后传入run_insert的，测试时需确保Mock逻辑能生成对应值，或在参数验证时匹配该字段。

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