ZOZO Advent Calendar 2025 2日目の記事になります。今回はlocalの環境でCloud Composer(Airflow)のテストを実行する方法を紹介したいと思います。

前準備

本題に入る前にDAGを実行する環境構築を行います。

環境構築

composer-local-devを用いて検証環境を構築します。コマンドのインストール方法は割愛します。

下記コマンドでdev環境の構築を行います。

$ composer-dev create \
  --from-image-version composer-2.15.4-airflow-2.9.3  \
  test --debug

composer-dev start test

上記コマンドでlocalのcomposer環境を立ち上げます。立ち上げに成功したらhttp://localhost:8080にアクセスしてairflowのUIが立ち上がったかを確認します

無事UIを確認できました。

DAGの作成

テスト対象のDAGを作成します。今回はDAGの実行時にパラメータを渡せるようにし、渡したパラメーターによって出力が変わるようなDAGを作成します。作成したDAGのコードは下記になります。

composer-test/composer/test/dags/tutorial.py

from datetime import timedelta
from airflow import DAG
from airflow.providers.standard.operators.python import (
    PythonOperator,
)
from airflow.models.param import Param


def get_username_from_ids(**kwargs):
    users = {
       1: {"username": "Yamada", "email": "user_1@example.com"},
       2: {"username": "Tanaka", "email": "user_2@example.com"},
       3: {"username": "Sato", "email": "user_3@example.com"},
       4: {"username": "Suzuki", "email": "user_4@example.com"},
    }
    user_ids = kwargs['params'].get("test_parameters").get("user_ids", [])
    return_user_info = []
    for user_id in user_ids:
      user = users.get(user_id)
      return_user_info.append(user)
    return return_user_info

default_args = {
    'owner': 'airflow',
    'depends_on_past': False,
    'email': ['airflow@example.com'],
    'email_on_failure': False,
    'email_on_retry': False,
    'retries': 1,
    'retry_delay': timedelta(minutes=5),

}
dag = DAG(
    'test DAG',
    default_args=default_args,
    description='Only use test DAG',
    tags=['example'],
    params={
        "test_parameters": Param(
            default={"key": "default_value"},
            type="object",
        )
    }
)

run_first = PythonOperator(task_id="print_the_context", python_callable=get_username_from_ids, dag=dag)

run_first

上記のDAGを手動実行する場合は、下記スクリーンショットの順でパラメーターを入力します。

今回の手動実行では下記jsonを入力します

{
  "user_ids": [1, 2]
}

実行結果のログを確認するとユーザー情報が返されています。

[2025-12-01, 19:01:49 UTC] {python.py:198} INFO - Done. Returned value was: [{'username': 'Yamada', 'email': 'user_1@example.com'}, {'username': 'Tanaka', 'email': 'user_2@example.com'}]

DAGのテスト

いよいよDAGのテストを実装します。

tests
├── expected
│   ├── test1.json
│   └── test2.json
├── sample_test.py
└── testdata
    ├── test1.json
    └── test2.json

上記のようなディレクトリをdagsディレクトリ配下に作成し、testdataディレクトリにinputを行うデータを配置します。expectedディレクトリに出力が期待されるデータを作成します。例としてtestdataディレクトリに下記2つのinputデータを作成しました。

testdata/test1.json

{
  "user_ids": [1, 2]
}

testdata/test2.json

{
  "user_ids": [3, 4]
}

expectedディレクトリには下記2つのjsonを配置します。

expected/test1.json

[
  {
    "username": "Yamada",
    "email": "user_1@example.com"
  },
  {
    "username": "Tanaka",
    "email": "user_2@example.com"
  }
]

expected/test2.json

[
  {
    "username": "Sato",
    "email": "user_3@example.com"
  },
  {
    "username": "Suzuki",
    "email": "user_4@example.com"
  }
]

最後にテストを行うpythonコードを実装します。expected, testdataディレクトリそれぞれを読み込み、メソッド経由でairflowのDAGを実行します。

sample_test.py

import pytest
from pathlib import Path
import yaml
import os
from airflow.models import DagBag, DagRun
testcases = [
  {"name": "テストケース1", "testcase_name": "test1"},
  {"name": "テストケース2", "testcase_name": "test2"},
]

def _load_testdata(testcase_name):
    testdata_path = Path(__file__).parent / "testdata" / f"{testcase_name}.json"
    
    with open(testdata_path, "r") as f:
        testdata = yaml.safe_load(f)
    return testdata

def _load_expected_result(testcase_name):
    expected_result_path = Path(__file__).parent / "expected" / f"{testcase_name}.json"
    with open(expected_result_path, "r") as f:
        expected_result = yaml.safe_load(f)
    return expected_result

def test_dag_response():
    dag_bag = DagBag(dag_folder=os.environ.get('DAGS_FOLDER'), include_examples=False)
    dag = dag_bag.get_dag(dag_id='test')
    assert dag is not None, "DAG 'tutorial' not found in DagBag"

    for testcase in testcases:
        testcase_name = testcase["testcase_name"]
        testdata = _load_testdata(testcase_name)

        dr = dag.test(run_conf={"test_parameters": testdata})
        ti = dr.get_task_instance(task_id="output_the_context")
        result = ti.xcom_pull()
        print("---- Test Result ----")

        expected_result = _load_expected_result(testcase_name)
        assert result == expected_result, f"Test case {testcase_name} failed: expected {expected_result}, got {result}"
        print(f"Testcase: {testcase_name}, Result: {result}" )

dag.test部分でDAGの実行を行い、get_task_instance部分でtaskの情報を取得します。取得した情報を出力が期待されるデータと突き合わせてassertを行っています。

テスト実行

試しに作成したテストが実行できるか確認します。下記コマンドを用いてコンテナ内に入ります。

docker exec -ti composer-local-dev-test /bin/bash

コンテナ内に入ったらpytestを実行します

$ pytest

実行を行うと無事テストが通ったことが確認できます

airflow@15756e2b1ac3:~$ pytest -v 
================================================================================= test session starts ==================================================================================
platform linux -- Python 3.11.8, pytest-7.4.4, pluggy-1.5.0 -- /usr/bin/python
cachedir: .pytest_cache
rootdir: /home/airflow
plugins: anyio-4.11.0
collected 1 item                                                                                                                                                                       

gcs/dags/tests/sample_test.py::test_dag_response PASSED                                                                                                                          [100%]

================================================================================== 1 passed in 1.40s ===================================================================================

まとめ

簡単にですがCloud Composerで用いるDAGのテストを行う方法をまとめました。WF系のツールはテストの作成が難しいと考えていましたが、Airflow側でテスト用のメソッドが用意されていたので楽にテストの実装ができました。

jongzの日記

Cloud Composer(Airflow)のDAGテストを実装してみる

前準備

環境構築

DAGの作成

DAGのテスト

テスト実行

まとめ