Airflow中PythonOperator/ShortCircuitOperator的Python函数单元测试及导入问题

解决导入失败问题 #

你遇到的导入失败，核心原因是Python默认路径未包含airflow根目录，导致tests目录下的脚本无法识别dags模块。有两种简单解决方式：

方式1：在测试脚本中添加路径 #

在test_my_dag.py开头加入以下代码，将airflow根目录加入Python搜索路径：

import sys
from pathlib import Path
# 获取airflow根目录（tests的父目录）
sys.path.append(str(Path(__file__).parent.parent))

方式2：用pytest配置文件统一设置 #

在airflow根目录创建pytest.ini文件，内容如下：

[pytest]
pythonpath = ./

之后直接在airflow根目录运行pytest，即可直接导入dags模块，无需每次手动加路径。

修正测试代码的Mock路径 #

你的Mock目标路径错误，Mock的路径需要和函数实际引用变量的位置一致。假设my_function在dags.my_dag中，且引用了同文件的my_variable，正确的Mock路径应为dags.my_dag.my_variable。

完整测试代码示例 #

假设dags/my_dag.py中的代码如下：

from airflow import DAG
from airflow.operators.python import ShortCircuitOperator
from datetime import datetime

my_variable = "original_value"

def my_function():
    # ShortCircuitOperator依赖此返回值决定是否执行后续任务
    return my_variable == "original_value"

with DAG(
    dag_id="my_dag",
    start_date=datetime(2024,1,1),
    schedule=None
) as dag:
    ShortCircuitOperator(
        task_id="check_condition",
        python_callable=my_function
    )

对应的tests/test_my_dag.py测试代码：

from unittest import mock
from dags.my_dag import my_function

@mock.patch('dags.my_dag.my_variable', "some_value")
def test_my_function():
    # 验证变量被Mock后，函数返回False
    assert my_function() is False

def test_my_function_default():
    # 验证默认状态下函数返回True
    assert my_function() is True

最佳实践建议 #

• 分离业务逻辑与DAG定义：把my_function这类业务判断函数从DAG文件中抽离，放到airflow/utils或airflow/operators目录下，单独维护和测试，避免和DAG配置代码耦合。
• 聚焦函数本身测试：ShortCircuitOperator/PythonOperator的核心是执行传入的Python函数，只要函数逻辑测试通过，Operator的行为就有保障，无需额外测试Operator本身（Airflow官方已覆盖Operator的基础逻辑测试）。

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