如何在pytest中对条件导入的模块进行打桩（Patch）？

这个问题我之前也碰到过，条件导入的对象在测试环境里找不到确实挺闹心的——核心原因就是你本地环境不满足Databricks的运行条件，导致业务模块里根本没生成F这个属性，打桩时自然会抛出找不到属性的错误。下面给你几个实用的解决办法：

方法一：提前给F设置占位符 #

在你的业务模块开头，先初始化一个F的占位符，不管环境是否满足导入条件，让这个属性始终存在：

import os

# 先初始化占位符，确保模块里始终有F这个属性
F = None

if "DATABRICKS_RUNTIME_VERSION" in os.environ:
    from pyspark.sql import functions as F

def save_results_to_databricks(...):
    # 正常使用F.col("relevant")等方法
    pass

这样本地环境下F是None，但模块里确实存在这个属性，测试时打桩就能正常找到它了。

方法二：在测试用例中提前注入F #

如果不想改动业务代码，可以在测试的前置步骤里，手动给目标模块添加F属性，再进行打桩：

import unittest
from unittest.mock import patch, Mock
import path.to.module  # 导入你的业务模块

class TestSaveResultsToDatabricks(unittest.TestCase):
    def setUp(self):
        # 检查模块是否有F属性，没有就先创建一个Mock占位
        if not hasattr(path.to.module, 'F'):
            path.to.module.F = Mock()
    
    @patch("path.to.module.F")
    def test_save_results_to_databricks(self, MockFunctions):
        # 模拟F.col的返回值
        MockFunctions.col.return_value = "mock_column"
        # 调用被测试函数
        path.to.module.save_results_to_databricks(...)
        # 断言方法被正确调用
        MockFunctions.col.assert_called_once_with("relevant")

方法三：用try-except导入并默认生成Mock #

这种方法更彻底，不管是Databricks还是本地环境，F都会存在，本地环境下直接用Mock代替：

try:
    # Databricks环境下正常导入
    from pyspark.sql import functions as F
except ImportError:
    # 本地环境导入失败，直接创建Mock对象
    from unittest.mock import Mock
    F = Mock()

def save_results_to_databricks(...):
    # 正常使用F的方法即可
    pass

这样本地开发和测试时，F本身就是Mock对象，你甚至可以直接在测试里配置它的行为，不需要额外打桩也能完成测试。

备注：内容来源于stack exchange，提问作者Julian