Pytest如何对依赖Snowflake连接返回DataFrame的函数做单元测试

需求可行性说明 #

完全可以实现，无需连接真实Snowflake数据库即可完成单元测试，同时保证业务代码被覆盖率统计命中。

原方案问题原因 #

你当前的做法是直接mock了extract_data_from_db这个属性本身，导致业务代码中该属性的实现逻辑完全没有被执行，所以不会被纳入覆盖率统计，相当于跳过了要测试的核心代码。

修正方案 #

核心思路是不mock目标测试方法本身，只mock它依赖的外部资源：即Snowflake的连接方法、以及pandas的SQL查询方法，这样既不会触发真实数据库请求，又能走完原方法的全部逻辑。

修正后的测试代码示例 #

假设你的业务代码存放在biz_module.py文件中，测试代码如下：

import pandas as pd
import pytest
from biz_module import Something

# 预设的mock返回数据
mock_df = pd.DataFrame({'a': [1]})

@pytest.fixture
def mock_external_deps(mocker):
    # mock snowflake的连接方法，避免真实建连
    mock_snow_conn = mocker.patch("biz_module.snowflake.connector.connect")
    # mock pandas的read_sql_query方法，返回预设数据，不触发真实查询
    mock_read_sql = mocker.patch("biz_module.pd.read_sql_query", return_value=mock_df)
    # 返回mock对象方便测试时做断言校验
    return mock_snow_conn, mock_read_sql

def test_extract_data_from_db(mock_external_deps):
    mock_snow_conn, mock_read_sql = mock_external_deps
    # 实例化对象，此时__init__中的connect是被mock的，不会真的连库
    instance = Something()
    # 调用属性方法，走原业务逻辑
    result = instance.extract_data_from_db
    
    # 断言返回值符合预期
    pd.testing.assert_frame_equal(result, mock_df)
    # 可选断言：校验SQL语句是否符合预期、连接参数是否正确
    mock_read_sql.assert_called_once_with('select * from table', mock_snow_conn.return_value)
    mock_snow_conn.assert_called_once_with(user='test_user', account='test_account', password='test_pwd')

效果说明 #

• 整个测试过程不会产生任何真实的Snowflake连接请求和数据库查询
• 原业务类的__init__方法、extract_data_from_db属性的全部逻辑都会被执行，会被正常纳入覆盖率统计
• 还可以额外校验SQL语句、连接参数的正确性，测试覆盖更全面

内容的提问来源于stack exchange，提问作者Ang Yiwei