如何用pytest测试内部调用Snowflake连接的类函数？

问题背景 #

需要测试SnowFlakeHelperClass中的to_be_tested_function函数，但该函数内部调用了_get_snowflake_connection方法建立实际数据库连接，导致测试无法正常运行。

类结构如下：

class SnowFlakeHelperClass:
    ...
    def _get_snowflake_connection(self) -> SnowflakeConnection:
       ...
    def to_be_tested_function(self):
       conn = self._get_snowflake_connection()
       # 后续使用conn的业务逻辑

现有测试代码：

def test_my_function(fixture_for_class):
   assert fixture_for_class.to_be_tested_function() == something

解决方法 #

核心思路是Mock掉_get_snowflake_connection方法，让它返回一个模拟的SnowflakeConnection对象，彻底规避实际数据库连接操作。

方法1：使用unittest.mock直接替换类方法 #

from unittest.mock import Mock, patch
import pytest

def test_my_function(fixture_for_class):
    # 创建模拟的数据库连接对象，spec参数保证模拟对象匹配原类结构
    mock_conn = Mock(spec=SnowflakeConnection)
    # 临时替换类实例的_get_snowflake_connection方法
    with patch.object(fixture_for_class, '_get_snowflake_connection', return_value=mock_conn):
        # 执行待测试函数
        result = fixture_for_class.to_be_tested_function()
        # 断言结果是否符合预期
        assert result == something
        # 可选：验证连接方法是否被正确调用
        fixture_for_class._get_snowflake_connection.assert_called_once()

方法2：在Fixture中提前Mock（适用于多测试场景） #

如果多个测试用例都需要规避数据库连接，可以直接在fixture中完成Mock配置：

import pytest
from unittest.mock import Mock

@pytest.fixture
def fixture_for_class():
    helper = SnowFlakeHelperClass()
    # 初始化模拟连接并替换原方法
    mock_conn = Mock(spec=SnowflakeConnection)
    helper._get_snowflake_connection = Mock(return_value=mock_conn)
    yield helper

def test_my_function(fixture_for_class):
    result = fixture_for_class.to_be_tested_function()
    assert result == something
    # 验证连接方法的调用情况
    fixture_for_class._get_snowflake_connection.assert_called_once()

方法3：使用pytest-mock插件简化代码 #

如果安装了pytest-mock插件，可进一步简化Mock操作：

def test_my_function(fixture_for_class, mocker):
    mock_conn = mocker.Mock(spec=SnowflakeConnection)
    # 一键替换目标方法
    mocker.patch.object(fixture_for_class, '_get_snowflake_connection', return_value=mock_conn)
    
    result = fixture_for_class.to_be_tested_function()
    assert result == something
    fixture_for_class._get_snowflake_connection.assert_called_once()

补充说明 #

• 如果to_be_tested_function中对conn有具体方法调用（比如conn.execute()），可以给mock_conn设置对应的返回值，例如mock_conn.execute.return_value = 预期结果，确保测试覆盖完整逻辑。
• 使用spec=SnowflakeConnection能避免测试中调用原类不存在的方法，保证Mock对象的合理性。

内容的提问来源于stack exchange，提问作者ASHISH M.G