How to Unit Test an Airflow DAG?

airflow

Developing Airflow dags involves writing unit tests for the individual tasks, and then manually running the whole dag from start to finish.

Here’s a simple operator for testing:

class MyOperator(BaseOperator):
    def execute(self, context):
        return 'foo'

To test the operator, first instantiate three objects:

  1. a DAG,
  2. a task (the operator we’re testing), and
  3. a TaskInstance.

Then call the execute method.

The Test

class TestMyOperator(TestCase):
    def test_execute(self):
        dag = DAG(dag_id='foo', start_date=datetime.now())
        task = MyOperator(dag=dag, task_id='foo')
        ti = TaskInstance(task=task, execution_date=datetime.now())
        task.execute(ti.get_template_context())

Or in Pytest:

def test_my_operator():
    dag = DAG(dag_id='foo', start_date=datetime.now())
    task = MyOperator(dag=dag, task_id='foo')
    ti = TaskInstance(task=task, execution_date=datetime.now())
    task.execute(ti.get_template_context())

Be sure to hide the globals section in definition files.

Running the whole dag

See how I run an entire dag from the commandline and watch the logs in real-time.