Testing is an integral part of any software system to build confidence and increase the reliability of the system. This deck covered different categories of tests that you can write for airflow. It includes DAG validation tests, pipeline definition tests, unit tests. Also, I will showcase all these tests using some examples.

1. Testing in Airflow
Chandu Kavar
Data Engineer
Grab, Ex-ThoughtWorker
@chandukavar

2. Traditional Web App Testing
We mostly write -
● Unit Tests
● Integration Tests
● Functional Tests
Client
Database
Server
Stress testing, snapshot testing etc. for
complex and large scale application.

3. ● Typos, cyclicity, mandatory parameters etc. in the DAGs
● No of tasks, nature of tasks in the DAG
● Upstream and downstream dependencies of each task
● Custom Operator, sensor etc.
● Communication between tasks such as xComs
● End-to-End pipeline
What can we test in Airflow?

4. Five Categories of tests we can write -
1. DAG validation tests
2. DAG/Pipeline definition tests
3. Unit tests for custom operators, sensor etc.
4. Integration tests
5. End-to-End pipeline tests

5. Sample DAG
dag = DAG('hello_world',
description='Hello world example',
schedule_interval='0 12 * * *',
start_date=datetime(2017, 3, 20))
dummy_operator = DummyOperator(task_id='dummy_task', retries = 3, dag=dag)
hello_operator = PythonOperator(task_id='hello_task', python_callable=print_hello, dag=dag)
multiplyby5_operator = MultiplyBy5Operator(my_operator_param= 20, task_id='multiplyby5_task',
dag=dag)
dummy_operator >> hello_operator
dummy_operator >> multiplyby5_operator
def print_hello():
return 'Hello'

6. DAG Validation Tests
class TestDagIntegrity(unittest.TestCase):
def setUp(self):
self.dagbag = DagBag()
def test_import_dags(self):
self.assertFalse(
len(self.dagbag.import_errors),
'DAG import failures. Errors: {}'.format(
self.dagbag.import_errors
)
)
def test_queue_present(self):
for dag_id, dag in self.dagbag.dags.iteritems():
queue = dag.default_args.get('queue', None)
msg = Queue not set for DAG {id}'.format(id=dag_id)
self.assertNotEqual(None, queue, msg)

7. DAG Definition Tests
class TestHelloWorldDAG(unittest.TestCase):
"""Check HelloWorldDAG expectation"""
def setUp(self):
self.dagbag = DagBag()
def test_task_count(self):
"""Check task count of hello_world dag"""
dag_id='hello_world'
dag = self.dagbag.get_dag(dag_id)
self.assertEqual(len(dag.tasks), 3)
def test_contain_tasks(self):
"""Check task contains in hello_world dag"""
dag_id='hello_world'
dag = self.dagbag.get_dag(dag_id)
tasks = dag.tasks
task_ids = list(map(lambda task: task.task_id, tasks))
self.assertListEqual(task_ids, ['dummy_task', 'multiplyby5_task','hello_task'])

8. DAG Definition Tests
def test_dependencies_of_dummy_task(self):
"""Check the task dependencies of dummy_task in hello_world dag"""
dag_id='hello_world'
dag = self.dagbag.get_dag(dag_id)
dummy_task = dag.get_task('dummy_task')
upstream_task_ids = list(map(lambda task: task.task_id, dummy_task.upstream_list))
self.assertListEqual(upstream_task_ids, [])
downstream_task_ids = list(map(lambda task: task.task_id, dummy_task.downstream_list))
self.assertListEqual(downstream_task_ids, ['hello_task', 'multiplyby5_task'])

9. Custom Operator
class MultiplyBy5Operator(BaseOperator):
@apply_defaults
def __init__(self, my_operator_param, *args, **kwargs):
self.operator_param = my_operator_param
super(MultiplyBy5Operator, self).__init__(*args, **kwargs)
def execute(self, context):
log.info('operator_param: %s', self.operator_param)
return (self.operator_param * 5)
class MultiplyBy5Plugin(AirflowPlugin):
name = "multiplyby5_plugin"

10. Unit tests - Custom Operator
class TestMultiplyBy5Operator(unittest.TestCase):
def test_execute(self):
dag = DAG(dag_id='anydag', start_date=datetime.now())
task = MultiplyBy5Operator(my_operator_param=10, dag=dag, task_id='anytask')
ti = TaskInstance(task=task, execution_date=datetime.now())
result = task.execute(ti.get_template_context())
self.assertEqual(result, 50)

11. https://github.com/chandulal/airflow-testing
Don’t forget to give a star!

12. Blog post on medium -
https://medium.com/@chandukavar

13. Thanks!