Open2

Airflow 2.6.0 で unittestを書く

Masato NakamuraMasato Nakamura

まずは1ファイルにテスト(unittest)もDAGも両方書いてみる

import unittest
from datetime import datetime
from airflow import DAG
from airflow.operators.bash import BashOperator

def create_dag():
    default_args = {
        'owner': 'airflow',
        'start_date': datetime(2023, 5, 6)
    }

    dag = DAG(
        'hello_world_dag',
        default_args=default_args,
        schedule_interval='@daily',
        catchup=False
    )

    echo_hello_world_task = BashOperator(
        task_id='echo_hello_world',
        bash_command='echo "Hello, World!"',
        dag=dag
    )

    return dag

class TestHelloWorldDag(unittest.TestCase):

    def setUp(self):
        self.dag = create_dag()

    def test_task_count(self):
        """Test if the number of tasks in the DAG is correct."""
        self.assertEqual(len(self.dag.tasks), 1)

    def test_bash_operator_task(self):
        """Test if BashOperator has the correct settings."""
        echo_hello_world_task = self.dag.get_task('echo_hello_world')

        self.assertIsInstance(echo_hello_world_task, BashOperator)
        self.assertEqual(echo_hello_world_task.bash_command, 'echo "Hello, World!"')

if __name__ == '__main__':
    unittest.main()

airflow breezeで実行してみる

$ breeze start-airflow # breeze環境を実行 TODO: 後でかく
$ pytest /files/dags/hello_world_dag.py

無事実行成功