這篇來講一下在測試的時候遇到的問題，主要是 package 的 import issue

一般來說，我們開發的 dag 都會放在 project 下的 /dags 裡面

project/
  dags/
    example_dag1.py
    example_dag2.py

如果你的專案夠大，你可能會想抽出一些共用的程式碼，放在 common 內之類的，假設我們放在 dags/common.py 內，並在 example_dag1.py import 它

project/
  dags/
    example_dag1.py
    example_dag2.py
    common.py

# example_dag1.py
from common import tool_function

dag = DAG(...)

可以發現一個跟一般常見 python project 不同的部份，我們竟然是直接 import common，而不是 from dags.common

或許你會覺得，因為在同一層所以用相對路徑也是可行的，但其實這是因為 Airflow 有將 dags/ 放到 PYTHONPATH 內，如果你真的寫成 from dags.common ，反而會在 Web UI 上出現錯誤而無法執行。

在這個階段，其實還不會造成測試上的錯誤，因為在同一層內，所以呼叫上是可以的。但如果我們的專案更加複雜了，我們決定將一些 PythonOperator 的 callable function 再抽出來，並放到下一層的 package 內：

project/
  dags/
    code/
      init.py
      execute_code1.py
    example_dag1.py
    example_dag2.py
    common.py

在 dags/comde/execute_cod1.py 內，我們也需要用到 common 檔

# dags/comde/execute_cod1.py
from dags.common import tool_function

def a_etl_step_function():
    # do something
    return 1

恭喜你，你的程式就會出錯了，原理同上。當然也一樣，你改回 from common import tool_function 這樣就可以過了，儘管你的 IDE 有很高的機率會跟你說這個路徑不正確，但是在 Airflow 的 WebUI 介面上是可以過的。

但是，假如我們接下來要對 execute_code1.py 寫單元測試的話？

# test_execute_code1.py
from dags.code.execute_code1 import a_etl_step_function

def test_etl_step_func():
  result = a_etl_step_function()
  assert result == 1

啊，恭喜你，pytest 十之八九會報 import error。因為 pytest 可不知道 Airflow 背後做了什麼手腳，對它來說 from common import tool_function 就是一個不正確的路徑。

該怎麼辦？手動加囉。

export PYTHONPATH="${PYTHONPATH}:${AIRFLOW_HOME}/dags"

AIRFLOW_HOME為你的專案路徑，如果你是 linux，那通常是 ~/airflow。

如此一來，Airflow 在執行時就也能讀到 dags 的這層，於是你的 import 就可以改回正確的 from dags.common import tool_function dags 開頭，那 pytest 也就能正確執行了。

這個問題困擾了我在開發測試時好久，因為實在不想把所有 py 都放在同一層內，但分到不同 package 內卻又會造成測試失敗，分享出來共勉之。