Airflow

Dag Loader Test - 로딩 오류와 파싱 시간 확인

Airflow 사용자는 Dag를 프로덕션 레벨 코드로 취급해야 하며, Dag가 기대한 결과를 내는지 확인하는 다양한 연관 테스트를 갖춰야 한다.

Dag Loader Test는 Dag를 로딩할 때 에러를 던지는 코드가 없는지 확인하는 테스트다. 사용자가 추가로 코드를 작성할 필요는 없다 - ``python your-dag-file.py``를 실행해 에러 없이 끝나면 미설치 의존성, 문법 오류 등이 없다는 것을 확인할 수 있다. 다만 스케줄러 환경과 동일한 의존성, 환경 변수, Dag가 참조하는 공통 코드를 갖춘 환경에서 로드해야 한다.

이 방법은 최적화 이후 Dag가 더 빠르게 로드되는지 확인하는 데도 좋다. Dag를 실행하고 걸리는 시간을 측정하되, 역시 동일한 의존성·환경 변수·공통 코드 조건에서 실행해야 한다. Linux에서는 내장 ``time`` 명령으로 측정할 수 있으며, 캐싱 효과를 감안해 여러 번 연속 실행해봐야 한다. 최적화 전후를 동일 조건(동일 머신, 환경 등)에서 비교해 효과를 판단한다.

``time python airflow/example_dags/example_python_operator.py``를 실행하면 real/user/sys 세 가지 시간이 나온다. 중요한 지표는 'real time'으로, Dag를 처리하는 데 걸린 실제 시간을 말한다. 이 방식으로 파일을 로드하면 새 인터프리터를 매번 시작하므로, Airflow가 Dag를 파싱할 때는 없는 초기 로딩 시간이 포함된다는 점에 유의해야 한다. 초기화 시간은 ``time python -c ''``로 별도 측정할 수 있다(예: real 0.073s) - 예제의 경우 인터프리터 시작 시간이 전체 파싱 시간(약 0.699s)의 약 10%를 차지하므로, 실제 파싱 시간은 약 0.62s로 추정할 수 있다.

단위 테스트 - Dag 로딩과 구조

단위 테스트는 Dag에 잘못된 코드가 없음을 보장한다. 태스크와 Dag 양쪽 모두에 대해 단위 테스트를 작성할 수 있다.

Dag 로딩에 대한 단위 테스트는 ``DagBag()`` fixture를 만들고, ``dagbag.get_dag(dag_id="hello_world")``로 Dag를 가져와 ``dagbag.import_errors == {}``, ``dag is not None``, ``len(dag.tasks) == 1`` 등을 assert하는 방식으로 작성한다.

Dag 구조에 대한 단위 테스트는, 코드로 생성된 Dag의 구조를 dict 객체와 비교하는 방식으로 작성할 수 있다. 예를 들어 ``assert_dag_dict_equal`` 헬퍼가 ``dag.task_dict.keys() == source.keys()``, 각 태스크 존재 여부(``dag.has_task``), 다운스트림 태스크 집합(``task.downstream_task_ids == set(downstream_list)``)을 비교하는 식이다.

커스텀 오퍼레이터/센서/deferrable 오퍼레이터 단위 테스트

커스텀 오퍼레이터나 센서를 단위 테스트하려면 직접 호출하면 된다. Dag Run, Dag, 메타데이터 데이터베이스가 필요하지 않다 - 오퍼레이터를 인스턴스화하고, 코드가 실제로 읽는 context 키만 넣어 ``execute()``를 호출하면 된다. 예를 들어 ``op = MyCustomOperator(task_id=..., prefix="s3://bucket/some/prefix")``를 만든 뒤 ``op.execute(context={}) == "expected return value"``를 assert한다.

센서는 ``poke()``를 호출해 반환된 boolean을 assert한다: ``sensor.poke(context={}) is True``.

오퍼레이터가 템플릿 필드를 렌더링한다면, assert 하기 전에 먼저 템플릿을 렌더링해야 한다: ``op.render_template_fields(context={"ds": "2021-09-13"})``를 호출한 뒤 ``op.prefix == "s3://bucket/2021-09-13"``를 확인한다.

deferrable 오퍼레이터는, 기대한 트리거로 defer하는지 확인한 다음 resume 경로를 직접 구동해 테스트한다. ``pytest.raises(TaskDeferred)``(``airflow.sdk.exceptions``에서 import)로 감싸 ``op.execute(context={})``를 호출하고, ``exc.value.trigger``가 기대한 트리거 클래스의 인스턴스인지 확인한 뒤, ``getattr(op, exc.value.method_name)(context={}, event={"status": "success"})``처럼 트리거가 resume하는 메서드를 직접 호출해 검증한다.

.. note::

``TaskInstance.run()``과 ``TaskInstance.render_templates()``는 Airflow 3.2에서 제거되었다 - ``TaskInstance``는 Airflow 3.0부터 내부(internal) 클래스로 취급된다. 기존에 ``ti.run()``을 쓰던 테스트 코드는 ``op.execute(context)``로, ``ti.render_templates(context)``는 ``op.render_template_fields(context)``로 대체해야 한다.

Dag Run 전체를 실행해보는 테스트를 하려면 ``dag.test()``를 사용한다. 이는 통합 테스트(integration test)에 해당한다 - 메타데이터 데이터베이스와, Airflow가 직렬화할 수 있는 Dag가 필요하므로 위에서 설명한 단위 테스트를 대체하는 수단은 아니다.

Self-Checks와 스테이징 환경

Dag 안에 태스크가 기대한 결과를 만들어내는지 확인하는 체크를 구현할 수도 있다. 예를 들어 S3에 데이터를 push하는 태스크가 있다면, 다운스트림 태스크에서 그 파티션이 S3에 실제로 생성됐는지, 데이터가 올바른지 간단히 확인하는 체크를 둘 수 있다.

마찬가지로 Kubernetes나 Mesos에서 마이크로서비스를 시작하는 태스크가 있다면, :class:`~airflow.providers.http.sensors.http.HttpSensor`로 그 서비스가 시작됐는지 확인해야 한다. 예시로 ``PushToS3`` 태스크 뒤에 ``S3KeySensor``로 ``s3://bucket/key/foo.parquet``가 존재하는지 확인하는 체크 태스크(``poke_interval=0, timeout=0``)를 이어붙일 수 있다.

가능하다면 프로덕션에 배포하기 전 전체 Dag 실행을 테스트할 스테이징 환경을 두는 것이 좋다. Dag는 변수를 바꿀 수 있도록 파라미터화해야 한다(예: S3 작업의 출력 경로나 설정을 읽어올 데이터베이스). Dag 안에 값을 하드코딩하고 환경에 따라 수동으로 바꾸면 안 된다. 대신 환경 변수를 이용해 Dag를 파라미터화할 수 있다 - 예: ``dest = os.environ.get("MY_DAG_DEST_PATH", "s3://default-target/path/")``.

Variable/Connection을 DB 없이 모킹하기

Variable이나 Connection을 사용하는 코드를 테스트할 때는 테스트 실행 시점에 그 객체들이 존재해야 한다. 가장 직관적인 해법은 이 객체들을 DB에 저장해 코드 실행 중 읽히게 하는 것이지만, DB에 객체를 읽고 쓰는 데는 추가적인 시간 오버헤드가 든다. 테스트 실행 속도를 높이려면, DB에 저장하지 않고 이 객체들의 존재를 시뮬레이션하는 것이 낫다. 이를 위해 :meth:`unittest.mock.patch.dict`로 :data:`os.environ`을 모킹해 환경 변수를 만들 수 있다.

Variable에는 :envvar:`AIRFLOW_VAR_{KEY}` 형식을 사용한다: ``with mock.patch.dict("os.environ", AIRFLOW_VAR_KEY="env-value"): assert "env-value" == Variable.get("key")``.

Connection에는 :envvar:`AIRFLOW_CONN_{CONN_ID}` 형식을 사용한다. ``Connection(conn_type="gcpssh", login="cat", host="conn-host")``을 만들고 ``conn.get_uri()``로 URI를 얻은 뒤, ``mock.patch.dict("os.environ", AIRFLOW_CONN_MY_CONN=conn_uri)``로 모킹하면 ``Connection.get("my_conn").login``이 ``"cat"``이 되는지 확인할 수 있다.