Pytest fixture模拟airflow.get_current_context触发AirflowException失败排查

问题描述 #

我尝试用Pytest的mocker fixture模拟airflow.operators.python.get_current_context，代码如下：

@pytest.fixture
def _mock_get_current_context(mocker):
    mocker.patch(
        "airflow.operators.python.get_current_context",
        return_value={},
    )

该模拟方式对requests.get、airflow.operators.trigger_dagrun.TriggerDagRunOperator.execute等其他函数均有效，但在DAG任务中调用get_current_context()时，仍触发如下错误：

if not _CURRENT_CONTEXT:
    raise AirflowException(
        "Current context was requested but no context was found! "
        "Are you running within an airflow task?"
        )
E       airflow.exceptions.AirflowException: Current context was requested but no context was found! Are you running within an airflow task?

查看get_current_context()源码，确认模拟未生效：

def get_current_context() -> Context:
    if not _CURRENT_CONTEXT:
        raise AirflowException(
            "Current context was requested but no context was found! "
            "Are you running within an airflow task?"
        )
    return _CURRENT_CONTEXT[-1]

模拟失败的原因及解决方法 #

1. 补丁路径错误（最常见） #

Airflow的get_current_context实际定义在airflow.utils.context模块，airflow.operators.python只是做了导入转发。如果你的DAG任务代码最终调用的是原始定义的函数，而你补丁的是转发模块的引用，就会导致mock失效。

另外，mock的核心是替换被测试代码中实际引用的对象，而非原始模块中的对象。比如如果你的DAG代码是from airflow.operators.python import get_current_context，那么你需要补丁的是你的_dag模块.get_current_context，而不是airflow.operators.python.get_current_context。

解决方法：

• 直接补丁原始定义的路径：

  @pytest.fixture
  def _mock_get_current_context(mocker):
      mocker.patch(
          "airflow.utils.context.get_current_context",
          return_value={},
      )
  

• 或者补丁被测试代码中导入的引用（假设DAG在dags/my_dag.py）：

  @pytest.fixture
  def _mock_get_current_context(mocker):
      mocker.patch(
          "dags.my_dag.get_current_context",
          return_value={},
      )
  

2. Fixture未被正确激活 #

如果测试函数没有声明使用_mock_get_current_context fixture，mock不会被应用到测试流程中。

解决方法：
在测试函数参数中声明该fixture：

def test_dag_task(_mock_get_current_context):
    # 执行测试逻辑
    ...

3. 导入顺序问题 #

如果被测试的DAG模块在mock补丁之前就已被导入，那么模块中已加载的get_current_context引用还是原始函数，不会被后续的mock替换。

解决方法：
调整顺序，先执行mock补丁，再导入DAG模块：

def test_dag_task(mocker):
    # 先补丁
    mocker.patch("airflow.utils.context.get_current_context", return_value={})
    # 再导入DAG模块
    from dags import my_dag
    # 执行测试逻辑
    ...

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