Airflow PyTest中DagBag导入本地模块_base.py失败问题

问题场景 #

我的DAG目录结构如下：

.
├── __init__.py
├── _base.py
├── dag1.py
├── dag2.py
├── dag3.py
├── ...

在Google Cloud Composer部署环境中，DAG文件可正常导入_base.py，但本地运行PyTest测试DagBag时，所有DAG文件均报以下错误：

ERROR - Failed to import: dags/dag2.py
Traceback (most recent call last):
  File "~/code/myproject/venv/lib/python3.7/site-packages/airflow/models/dagbag.py", line 339, in parse
    loader.exec_module(new_module)
  File "<frozen importlib._bootstrap_external>", line 728, in exec_module
  File "<frozen importlib._bootstrap>", line 219, in _call_with_frames_removed
  File "dags/dag2.py", line 12, in <module>
    from _base import start_here, task_factory
ModuleNotFoundError: No module named '_base'

相关代码片段 #

_base.py 通用任务辅助文件 #

def task_factory(
    operator_func,
    base_env_vars=[],
    op_kwargs={},
):
    def build_airflow_task(task_id=None, extra_env_vars=[], **kwargs):
        task = operator_func(
            task_id=task_id,
            name=task_id,
            env_vars=[*base_env_vars, *extra_env_vars],
            **op_kwargs,
            **kwargs,
        )
        return task

    return build_airflow_task

DAG文件中的导入代码 #

import os

from airflow import DAG
from airflow.kubernetes.secret import Secret
from airflow.providers.google.cloud.operators.kubernetes_engine import (
    GKEStartPodOperator,
)
from airflow.utils.helpers import chain
from kubernetes.client import models as k8s

from _base import start_here, finish, dummy_task, task_factory

测试文件代码 #

import sys

sys.path.insert(0, "..")

from airflow.models import DagBag


def test_dagbag_compiles():
    dags = DagBag("dags", include_examples=False)
    assert len(dags.import_errors) == 0

解决方案 #

方案1：调整测试文件的Python路径 #

修改测试文件，将DAG目录的绝对路径加入sys.path，确保Python能找到_base.py模块：

import os
import sys

# 将dags目录的绝对路径加入Python路径
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../dags")))

from airflow.models import DagBag


def test_dagbag_compiles():
    dags = DagBag("dags", include_examples=False)
    assert len(dags.import_errors) == 0

修改后，DagBag解析DAG文件时，Python会在dags目录中查找_base.py。

方案2：使用相对导入修改DAG文件 #

将DAG文件中的导入语句改为相对导入（适用于Python 3+）：

# 替换原来的绝对导入
from ._base import start_here, finish, dummy_task, task_factory

这种方式依赖于DAG目录作为Python包（已有__init__.py），Composer环境也支持这种导入方式，同时本地测试时也能正确解析。

方案3：设置PYTHONPATH环境变量 #

运行PyTest时，将DAG目录添加到PYTHONPATH中：

PYTHONPATH=./dags pytest tests/

这种方式不需要修改代码，通过环境变量告诉Python去哪里找模块。

内容的提问来源于stack exchange，提问作者K Pekosh