如何修复Snowflake数据获取函数的单元测试失败问题？

问题描述 #

编写的Snowflake数据获取Python函数单元测试断言失败，返回空列表[]，同时报错：

解密字符串时出错。必须指定区域。

报错原因 #

1. 测试未Mockdecrypt函数：原函数调用service.kms_utils.decrypt解密Snowflake凭证，但单元测试仅Mock了snowflake.connector.connect，实际执行时触发真实解密逻辑，因缺少区域参数抛出异常。
2. 异常被捕获后返回空列表：函数的try-except块捕获所有异常，最终返回初始定义的空列表server_details，导致断言失败。

修复步骤 #

1. Mock decrypt函数 #

避免测试时调用真实KMS解密逻辑，直接返回模拟凭证值。

2. （可选）Mock环境变量获取 #

若需模拟环境变量不存在场景，可Mockos.environ.get，否则直接让decrypt返回固定值即可。

3. 修复SQL注入风险（额外优化） #

原函数用字符串拼接生成SQL，存在注入风险，改为参数化查询更安全。

完整修复代码 #

修改后的业务函数（优化SQL） #

import os
import snowflake.connector 
from service.kms_utils import decrypt 
import logging

logger = logging.getLogger(__name__)

def get_data_from_snowflake(server_name):
    # 使用参数化查询避免SQL注入
    sql_query = "Select * from mytable where servername=%s"
    server_details = []
    try:
        conn = snowflake.connector.connect(
            user=decrypt(os.environ.get('sfusername')),
            password=decrypt(os.environ.get('sfpwd')),
            account=decrypt(os.environ.get('account')),
            warehouse="wh1"
        )
    
        cs = conn.cursor()
        # 传入参数执行查询
        cs.execute(sql_query, (server_name,))
        server_details = cs.fetchall()
        conn.close()
    
    except Exception as e:
        if "Could not connect to Snowflake Query backend after" in str(e):
            logger.error(f"Snowflake: Snowflake query timeout. {e}")
        elif "Failed to connect to DB" in str(e):
            logger.error(f"Snowflake: Snowflake DB connection error. {e}")
        else:
            logger.error(f"Snowflake: Failed to get CI validation data. {e}")
    return server_details

修复后的单元测试代码 #

import pytest
from unittest.mock import patch, MagicMock
from my_module import get_data_from_snowflake


@patch('snowflake.connector.connect')
@patch('my_module.decrypt')  # Mock业务函数中导入的decrypt函数
def test_get_data_from_snowflake_success(mock_decrypt, mock_connect):
    # 模拟decrypt返回固定凭证值
    mock_decrypt.return_value = "mock_value"
    
    # Mock Snowflake连接与游标逻辑
    mock_cursor = MagicMock()
    mock_connect.return_value.cursor.return_value = mock_cursor
    # 模拟查询结果
    mock_cursor.fetchall.return_value = [(1, 'Server1', 'Details1'), (2, 'Server2', 'Details2')]

    # 调用测试函数
    result = get_data_from_snowflake('Server1')

    # 断言结果符合预期
    assert result == [(1, 'Server1', 'Details1'), (2, 'Server2', 'Details2')]
    # 验证参数化查询是否传入正确参数
    mock_cursor.execute.assert_called_once_with("Select * from mytable where servername=%s", ('Server1',))

补充说明 #

• Mockdecrypt时，需注意patch路径为业务函数中导入decrypt的路径（即my_module.decrypt），而非service.kms_utils.decrypt，这是unittest.mock的路径规则。
• 若需测试异常场景，可让mock_decrypt抛出异常，验证logger是否正确记录错误。

内容的提问来源于stack exchange，提问作者Chinmay Nayak