调用函数时Mock Snowflake连接失效问题求助

问题背景 #

我编写了utils/snowflake_manager.py文件，其中SnowflakeManager类的format_some_input_data方法会调用get_table_type获取表类型（该方法实际连接Snowflake执行查询）。为测试format_some_input_data，我在tests/test_snowflake_manager.py中使用@patch装饰器尝试Mockget_table_type方法，使其返回"VIEW"，但运行Pytest时仍会连接真实Snowflake数据库，因测试表不存在导致table_type为None，测试失败。

原业务代码 #

# utils/snowflake_manager.py
import snowflake.connector

class SnowflakeManager:
    def __init__(self):
        self.conn = snowflake.connector.connect(
            user="user", .... , warehouse="warehouse")

    def get_table_type(self, table_name):
        query = f"""SELECT ... .. WHERE table_name = '{table_name}'"""

        table_or_view = self.conn.cursor().execute(query).fetchone()[0] # 应返回'VIEW'或'TABLE'

        return table_or_view

    def format_some_input_data(self, input):
        table_name = ... 对输入参数的转换逻辑
        table_type = self.get_table_type(table_name)

        final = ... 使用table_type值的后续转换逻辑

        return final

if __name__ == "__main__":
  snowflake_manager = SnowflakeManager()

原测试代码 #

from utils import SnowflakeManager # utils目录下有__init__文件直接导出SnowflakeManager

import pytest
from unittest.mock import patch

@patch("utils.snowflake_manager")
def test_get_table_type(mock_get_table_type):
    input_dummy_data = [...{"table": "someTableThatDoesntExistInSnowflake"}...]

    expected = [...{"table": "someTableThatDoesntExistInSnowflake", "table_type": "VIEW"},...]

    mock_get_table_type.get_table_type.return_value = "VIEW"
    
    assert SnowflakeManager().format_some_input_data(input_dummy_data) == expected

错误原因 #

• Patch路径错误：当前patch的是utils.snowflake_manager模块，而非SnowflakeManager类中的get_table_type方法。Mock的核心原则是patch对象被调用的位置，这里format_some_input_data是调用自身实例的get_table_type，所以需要patch类的方法而非整个模块。
• 初始化触发真实连接：SnowflakeManager的__init__方法直接创建了Snowflake连接，即使Mock了get_table_type，实例化类时仍会执行真实连接逻辑，导致测试时访问真实数据库。

修正后的测试代码 #

from utils import SnowflakeManager
import pytest
from unittest.mock import patch, MagicMock

@patch("utils.snowflake_manager.snowflake.connector.connect")
@patch("utils.SnowflakeManager.get_table_type")
def test_format_some_input_data(mock_get_table_type, mock_connect):
    # Mock get_table_type返回指定值
    mock_get_table_type.return_value = "VIEW"
    # Mock连接对象，避免初始化时建立真实连接
    mock_conn = MagicMock()
    mock_connect.return_value = mock_conn

    input_dummy_data = [{"table": "someTableThatDoesntExistInSnowflake"}]
    expected = [{"table": "someTableThatDoesntExistInSnowflake", "table_type": "VIEW"}]

    manager = SnowflakeManager()
    result = manager.format_some_input_data(input_dummy_data)
    
    assert result == expected
    # 验证get_table_type是否被正确调用
    mock_get_table_type.assert_called_once_with("someTableThatDoesntExistInSnowflake")

修正说明 #

1. 调整Patch路径：改为直接patchutils.SnowflakeManager.get_table_type，确保Mock的是类的目标方法
2. Mock数据库连接：添加对snowflake.connector.connect的Mock，阻止__init__中建立真实连接
3. 参数顺序正确：patch装饰器按从下往上的顺序生效，因此测试函数参数顺序为mock_get_table_type（最后一个patch）在前，mock_connect在后
4. 添加调用验证：通过assert_called_once_with确认get_table_type被正确调用，确保Mock逻辑生效

内容的提问来源于stack exchange，提问作者Jonathan Duran