在单元测试中为 Airflow 1.9 运行测试 Dag

Question

我已经实现了运行单个 dag 的测试用例，但它似乎在 1.9 中不起作用，可能是由于气流 1.8 中引入了更严格的池 . 我正在尝试在测试用例下运行：

from airflow import DAG
from airflow.operators.dummy_operator import DummyOperator
from airflow.operators.python_operator import PythonOperator

class DAGTest(unittest.TestCase):
    def make_tasks(self):
        dag = DAG('test_dag', description='a test',
                  schedule_interval='@once',
                  start_date=datetime(2018, 6, 26),
                  catchup=False)


        du1 = DummyOperator(task_id='dummy1', dag=dag)
        du2 = DummyOperator(task_id='dummy2', dag=dag)
        du3 = DummyOperator(task_id='dummy3', dag=dag)
        du1 >> du2 >> du3
        dag.run()

    def test_execute(self):
        self.make_tasks() 

例外：

Dependencies not met for <TaskInstance: test_dag.dummy3 2018-06-26 00:00:00  [upstream_failed]>, dependency 'Trigger Rule' FAILED: Task's trigger rule 'all_success' requires all
    upstream tasks to have succeeded, but found 1 non-success(es).
    upstream_tasks_state={'skipped': 0L, 'successes': 0L, 'failed': 0L,'upstream_failed': 1L, 'done': 1L, 'total': 1}, upstream_task_ids=['dummy2']

我做错了什么？ 我已经尝试过 LocalExecutor 和 SequentialExecutor

环境：

Python 2.7
Airflow 1.9

我相信它正在尝试同时执行所有任务而不尊重依赖关系。 注意：类似代码用于 Airflow 1.7

Answer 1

这是一个您可以在 pytest 测试用例中使用的函数，它将按顺序运行 DAG 的任务。

from datetime import timedelta
import pytest
from unittest import TestCase


@pytest.fixture
def test_dag(dag):
    dag._schedule_interval = timedelta(days=1)  # override cuz @once gets skipped
    done = set([])

    def run(key):
        task = dag.task_dict[key]
        for k in task._upstream_task_ids:
            run(k)
        if key not in done:
            print(f'running task {key}...')
            date = dag.default_args['start_date']
            task.run(date, date, ignore_ti_state=True)
            done.add(key)
    for k, _ in dag.task_dict.items():
        run(k)

然后您可以在测试中使用 test_dag(dag) 而不是 dag.run()。

您需要确保您在自定义运算符中的登录使用self.log.info() 而不是logging.info() 或print()，否则它们不会显示。

您可能还需要使用 python -m pytest -s test_my_dag.py 运行测试，因为没有 -s 标志将不会捕获 Airflow 的标准输出。

我仍在尝试弄清楚如何处理 DAG 间的依赖关系。

Answer 2

我不熟悉 Airflow 1.7，但我猜它没有 Airflow1.8 及更高版本所具有的“DagBag”概念。

你不能运行你这样创建的 DAG，因为dag.run() 启动了一个新的 python 进程，它必须从它在磁盘上解析的 dag 文件夹中找到 DAG - 但它不能。这作为消息包含在输出中（但您没有包含完整的错误消息/输出）

你想通过在测试文件中创建一个 dag 来测试什么？它是自定义运算符吗？那么你最好直接测试它。例如，这是我如何独立测试自定义运算符：

class MyPluginTest(unittest.TestCase)
    def setUp(self):
        dag = DAG(TEST_DAG_ID, schedule_interval='* * * * Thu', default_args={'start_date': DEFAULT_DATE})
        self.dag = dag
        self.op = myplugin.FindTriggerFileForExecutionPeriod(
            dag=dag,
            task_id='test',
            prefix='s3://bucket/some/prefix',
        )
        self.ti = TaskInstance(task=self.op, execution_date=DEFAULT_DATE)

        # Other S3 setup here, specific to my test


    def test_execute_no_trigger(self):
        with self.assertRaises(RuntimeError):
            self.ti.run(ignore_ti_state=True)

        # It shouldn't have anything in XCom
        self.assertEqual(
            self.ti.xcom_pull(task_ids=self.op.task_id),
            None
        )