How to Unit Test an Airflow DAG?
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:
- a DAG,
- a task (the operator we’re testing), and
- 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.