Airflow自定义插件编写与导入问题及测试最佳实践咨询

Alright, let's break down your problems step by step—you're hitting two key pain points with Airflow plugins: proper component registration/imports and efficient testing workflows. Here's how to fix both:

The root cause of your import issues is that Airflow doesn’t automatically scan and load all files in the plugins directory unless you explicitly register them via the AirflowPlugin class. This is the critical missing piece from your setup.

Where to Place the AirflowPlugin Class #

You need to define this class in the top-level __init__.py of your plugins directory—this is the only place Airflow looks for plugin registrations by default.

Example Registration Code #

Update your plugins/__init__.py with this code to register all your custom components:

from airflow.plugins_manager import AirflowPlugin
# Import all your custom hooks, operators, sensors, etc.
from hooks.my_hook import MyHook
from hooks.another_hook import AnotherHook
from operators.my_operator import MyOperator
from operators.another_operator import AnotherOperator

class CustomAirflowPlugins(AirflowPlugin):
    # Unique name for your plugin (must not conflict with other plugins)
    name = "custom_plugins"
    # List all hooks to register with Airflow
    hooks = [MyHook, AnotherHook]
    # List all operators to register with Airflow
    operators = [MyOperator, AnotherOperator]
    # Add sensors or utils here if you have them, e.g.:
    # sensors = [MyCustomSensor]
    # macros = [my_custom_utility_function]

Correct Import Paths #

Once registered, use these reliable import patterns:

• For internal plugin code (operators importing hooks, etc.): Use relative imports to avoid dependency on Airflow’s namespace. In my_operator.py, write:

  from ..hooks.my_hook import MyHook
  

• For DAG files outside the plugin: You can now import directly from Airflow’s core namespaces:

  from airflow.hooks.my_hook import MyHook
  from airflow.operators.my_operator import MyOperator
  

Post-Change Steps #

After updating, restart your Airflow scheduler, worker, and webserver. Double-check that your airflow.cfg has the correct plugins_folder set (default is AIRFLOW_HOME/plugins, which matches your structure).

You don’t need to modify Airflow configs every time you test—here’s how to streamline your workflow:

a. Avoid Dependency on Airflow’s Plugin Loader #

Treat your custom hooks/operators as regular Python modules for unit tests. Add your project root to the Python path so tests can import components directly.

If your project structure looks like this:

your_project/
├── plugins/
│   ├── hooks/
│   ├── operators/
│   └── __init__.py
└── tests/
    ├── test_my_operator.py
    └── test_my_hook.py

Add this to the top of your test files (or use a conftest.py in your project root for pytest to apply it globally):

import sys
from pathlib import Path

# Add project root to Python path
sys.path.append(str(Path(__file__).parent.parent))

# Now you can import directly
from plugins.operators.my_operator import MyOperator
from plugins.hooks.my_hook import MyHook

b. Mock Airflow Core Components #

Don’t spin up a full Airflow environment for unit tests—use mocking to simulate Airflow’s context objects (like ti, dag) or external dependencies. For example, using unittest.mock:

from unittest.mock import Mock, patch
from plugins.operators.my_operator import MyOperator

def test_my_operator_execute():
    # Mock the MyHook class and its methods
    with patch('plugins.operators.my_operator.MyHook') as mock_hook_cls:
        mock_hook = Mock()
        mock_hook_cls.return_value = mock_hook
        mock_hook.run_task.return_value = "success"

        # Create your operator instance
        op = MyOperator(task_id="test_task", target="test_target")
        # Execute the operator with a mock context
        result = op.execute(context={})

        # Assert the expected behavior
        assert result == "success"
        mock_hook.run_task.assert_called_once_with("test_target")

c. Integration Testing (Optional) #

If you want to verify your plugin loads correctly in Airflow, use temporary environment variables to point Airflow to your plugin directory without modifying airflow.cfg:

# List registered plugins to confirm your plugin is loaded
AIRFLOW__CORE__PLUGINS_FOLDER=/path/to/your/plugins airflow plugins list

# Test a single task from your DAG
AIRFLOW__CORE__PLUGINS_FOLDER=/path/to/your/plugins airflow tasks test your_dag_id your_task_id 2023-01-01

To make your plugin truly portable:

• Keep the plugins directory structure self-contained
• Use relative imports for all internal plugin code
• Include the AirflowPlugin registration in the top-level __init__.py
• When sharing, just tell users to copy the entire plugins folder into their AIRFLOW_HOME

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