如何Mock XCom以测试Airflow的Python可调用任务？

Airflow自带的DAG和任务测试工具无法覆盖仅测试单个Python可调用函数的场景，而这类函数中的XCom交互是测试的核心难点。以下是基于unittest和unittest.mock的解决方案，无需运行完整DAG或依赖Airflow元数据库。

核心思路 #

Python可调用函数中的XCom操作（xcom_pull/xcom_push）都是通过传入的context参数（尤其是其中的ti即Task Instance对象）实现的。只需Mock这个ti对象的XCom方法，就能模拟函数的XCom交互逻辑。

具体实现步骤 #

1. 准备被测函数示例 #

假设你有一个依赖XCom的处理函数：

def process_data(**context):
    # 从指定任务拉取原始数据
    raw_data = context['ti'].xcom_pull(task_ids='fetch_data_task')
    # 数据处理逻辑
    processed_data = [item * 2 for item in raw_data]
    # 将处理结果推回XCom
    context['ti'].xcom_push(key='processed_result', value=processed_data)
    return processed_data

2. 编写测试用例 #

使用unittest.mock.Mock来模拟Task Instance和XCom方法，验证函数逻辑和XCom调用行为：

import unittest
from unittest.mock import Mock
from your_module import process_data  # 替换为你的函数所在模块

class TestProcessData(unittest.TestCase):
    def test_process_data_xcom_interaction(self):
        # Mock Task Instance对象
        mock_ti = Mock()
        # 预设xcom_pull的返回值
        mock_ti.xcom_pull.return_value = [1, 2, 3]
        # 构造测试用的context
        test_context = {'ti': mock_ti}

        # 执行被测函数
        result = process_data(**test_context)

        # 断言数据处理结果正确
        self.assertEqual(result, [2, 4, 6])

        # 验证xcom_pull的调用参数符合预期
        mock_ti.xcom_pull.assert_called_once_with(task_ids='fetch_data_task')
        # 验证xcom_push的调用参数符合预期
        mock_ti.xcom_push.assert_called_once_with(key='processed_result', value=[2, 4, 6])

if __name__ == '__main__':
    unittest.main()

进阶技巧 #

• 模拟多XCom键值：如果函数需要拉取多个XCom键，可以用side_effect让xcom_pull根据参数返回不同值：

  def mock_xcom_pull(**kwargs):
      if kwargs.get('key') == 'raw_data':
          return [1, 2, 3]
      elif kwargs.get('key') == 'processing_config':
          return {'multiplier': 3}
  mock_ti.xcom_pull.side_effect = mock_xcom_pull
  

• 适配装饰器任务：对于用@task装饰器定义的函数，测试时只需将mock的context或ti传入函数即可，逻辑与普通函数一致。
• 验证调用次数：如果函数多次调用XCom方法，可以用assert_called_with、assert_called_times等方法验证调用行为。

内容的提问来源于stack exchange，提问作者S.aad