如何在单元测试中Mock从AWS SSM参数存储获取Airflow变量？

核心思路 #

直接Mock Airflow的Variable.get()方法即可，无需关心变量实际存储在SSM还是其他后端——因为Variable.get()是统一的上层调用入口，和底层存储实现解耦。

方法1：用unittest.mock直接Patch #

这是最直接的方式，针对你的变量键值对编写Mock逻辑：

from unittest.mock import patch
from airflow.models import Variable

def test_your_airflow_code():
    # 定义要Mock的SSM参数键值对
    mock_ssm_vars = {
        'a/b/c': 'true',
        'a/b/d': 'test'
    }

    # Patch Variable.get方法，根据传入的key返回对应值
    with patch.object(Variable, 'get') as mock_get:
        mock_get.side_effect = lambda key, **kwargs: mock_ssm_vars[key]

        # 调用你的业务代码（示例）
        key1 = Variable.get('a/b/c')
        key2 = Variable.get('a/b/d')

        # 验证结果和调用情况
        assert key1 == 'true'
        assert key2 == 'test'
        mock_get.assert_any_call('a/b/c')
        mock_get.assert_any_call('a/b/d')

方法2：用Pytest Fixture封装重复逻辑 #

如果多个测试用例需要Mock相同的变量，可以封装成Fixture复用：

import pytest
from unittest.mock import patch
from airflow.models import Variable

@pytest.fixture
def mock_ssm_variables():
    mock_vars = {
        'a/b/c': 'true',
        'a/b/d': 'test'
    }
    with patch.object(Variable, 'get') as mock_get:
        mock_get.side_effect = lambda key, **kwargs: mock_vars[key]
        yield mock_get

def test_task1(mock_ssm_variables):
    result = Variable.get('a/b/c')
    assert result == 'true'
    mock_ssm_variables.assert_called_once_with('a/b/c')

def test_task2(mock_ssm_variables):
    result = Variable.get('a/b/d')
    assert result == 'test'

为什么AIRFLOW_VAR_xxx无效？ #

Airflow的环境变量方式要求变量名是合法的环境变量格式（不能包含/等特殊字符），而你的变量路径a/b/c包含斜杠，无法转换成符合要求的环境变量名，因此这种方法不适用你的场景。

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