如何在单元测试中测试含多任务的Airflow DAG？

你遇到的情况很典型：单独调用task.run()只能执行单个任务，完全不会处理任务间的上下游依赖关系，所以多任务关联的DAG用这种方式没法按预期跑起来。下面给你两种实用的测试方案：

方案一：用Airflow CLI命令（推荐） #

Airflow自带的CLI工具专门针对DAG和任务测试做了优化，能完美处理依赖关系：

1. 测试完整DAG执行流程
   用airflow dags test命令，它会模拟触发一次完整的DAG运行，严格按照任务依赖顺序执行所有任务：

   airflow dags test your_dag_id 2024-01-01
   

   这里的2024-01-01是你指定的执行日期，这个命令会跳过调度器的常规逻辑，直接触发一次全量执行，非常适合快速验证依赖链是否正常。

2. 从首个任务开始执行（自动处理上游依赖）
   如果你只想从某个任务启动，同时确保它的上游依赖都被执行，可以用airflow tasks run命令加上--local参数（本地运行，不提交到调度器）：

   airflow tasks run your_dag_id task1 2024-01-01 --local
   

   这个命令会自动检查task1的所有上游任务，先执行完上游再运行task1，完全符合你的需求。

方案二：在代码中编写测试逻辑 #

如果你想在代码里直接触发测试，可以借助Airflow的TaskInstance类来处理依赖：

先补全你的DAG定义，再添加测试代码：

from airflow import DAG
from airflow.operators.bash_operator import BashOperator
from datetime import datetime, timedelta
from airflow.models import TaskInstance

# 定义DAG基础参数
default_args = {
    'owner': 'test_user',
    'start_date': datetime(2024, 1, 1),
    'retries': 1,
    'retry_delay': timedelta(minutes=5),
}

# 构建多任务DAG
with DAG(
    'multi_task_test_dag',
    default_args=default_args,
    schedule_interval=timedelta(days=1),
    catchup=False
) as dag:
    task1 = BashOperator(
        task_id='task1',
        bash_command='echo "Executing Task 1 successfully!"'
    )

    task2 = BashOperator(
        task_id='task2',
        bash_command='echo "Executing Task 2 after Task 1!"'
    )

    # 设置上下游依赖
    task1 >> task2

# 本地测试逻辑
if __name__ == "__main__":
    execution_date = datetime(2024, 1, 1)
    # 初始化task1的任务实例并运行，自动处理上游依赖（这里task1无上游，直接执行）
    ti_task1 = TaskInstance(task=task1, execution_date=execution_date)
    ti_task1.run(ignore_ti_state=True)
    
    # 运行task2时，会自动检查task1是否完成
    ti_task2 = TaskInstance(task=task2, execution_date=execution_date)
    ti_task2.run(ignore_ti_state=True)

这段代码里，TaskInstance.run()会自动识别任务的依赖关系，确保只有上游任务执行完成后，当前任务才会启动。ignore_ti_state=True是为了忽略任务实例的现有状态，强制触发执行，适合测试场景。

小提示 #

• 测试前确保Airflow元数据库已初始化，相关的连接、变量配置正确。
• 测试环境尽量用本地执行方式，避免干扰正式调度流程。

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