如何在导入Airflow DAG模块前Mock my_function函数？

核心逻辑是：Python导入模块时会执行顶层代码，因此必须在导入my_dag模块之前完成对my_function的Mock，避免触发实际函数的连接初始化逻辑。以下是两种常见测试框架下的实现方式：

方法一：使用标准库unittest.mock #

方式1：上下文管理器提前Mock #

# test_my_dag.py
from unittest.mock import patch

# 在导入my_dag前完成Mock
with patch('my_module.my_function') as mock_func:
    # 设置Mock返回值，需符合schedule_interval的格式要求
    mock_func.return_value = '@daily'
    # 此时导入my_dag，my_function已被Mock，不会执行实际逻辑
    from my_dag import MY_VAR

def test_my_var():
    assert MY_VAR == '你的预期值'

方式2：模块级别持久Mock #

# test_my_dag.py
from unittest.mock import patch

# 模块启动时立即Mock，确保导入my_dag前生效
mock_func = patch('my_module.my_function').start()
mock_func.return_value = '@daily'

# 导入目标变量
from my_dag import MY_VAR

def test_my_var():
    assert MY_VAR == '你的预期值'

# 测试结束后清理Mock
def teardown_module():
    patch.stopall()

方法二：使用pytest + pytest-mock #

方式1：自动生效的Fixture #

# test_my_dag.py
import pytest

@pytest.fixture(autouse=True, scope='session')
def mock_my_function(mocker):
    # 提前Mock目标函数
    mock_func = mocker.patch('my_module.my_function')
    mock_func.return_value = '@daily'
    # 在Fixture中完成导入，确保Mock生效
    global MY_VAR
    from my_dag import MY_VAR

def test_my_var():
    assert MY_VAR == '你的预期值'

方式2：直接前置Mock #

# test_my_dag.py
from pytest_mock import mocker

# 先Mock再导入
mocker.patch('my_module.my_function', return_value='@daily')
from my_dag import MY_VAR

def test_my_var():
    assert MY_VAR == '你的预期值'

关键注意事项 #

• Mock路径要准确：必须对应my_dag.py中导入my_function的实际路径，比如my_dag里是from my_module import my_function，则Mock路径为my_module.my_function。
• Mock返回值要合法：需符合Airflow schedule_interval的要求（如 cron 表达式、timedelta 实例或@daily这类特殊字符串），避免导入my_dag时因格式错误抛出异常。

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