如何在pytest导入Python脚本前模拟环境变量与依赖包？

我用pytest测试AWS Lambda函数时遇到了问题：Lambda代码（catalog_service/lambda_v1.py）在模块级别直接做了初始化操作——导入依赖后就读取os.environ['ENVIRONMENT']，调用boto3的SSM客户端拿参数，还初始化了AppSearch实例。

但运行测试时，导入模块的瞬间就会执行这些初始化代码，直接抛出KeyError说ENVIRONMENT不存在。我试过用monkeypatch模拟环境变量，但时机太晚，因为模块导入在测试代码执行之前就完成了。

怎么才能在导入Lambda模块前，先模拟好环境变量，同时mock掉boto3和AppSearch这些依赖？

Lambda代码示例（catalog_service/lambda_v1.py） #

import boto3
import botocore
import json
import os

from aws_lambda_powertools import Logger
from elastic_enterprise_search import AppSearch


LOGGER = Logger()
ENV = os.environ['ENVIRONMENT']
ssm = boto3.client('ssm')
# 注：原代码中key变量未定义，这里假设是根据ENV生成的路径
key = f"/{ENV}/app-search/api-key"
key_response = ssm.get_parameter(
    Name=key,
    WithDecryption=True
)
APP_SEARCH = AppSearch(
    os.environ['APP_SEARCH_API_HOST'],
    bearer_auth=key_response['Parameter']['Value'],
)

# 示例待测试函数
def format_good_response(data):
    return {
        'statusCode': '200',
        'headers': {'content-type': 'application/json'},
        'body': json.dumps(data)
    }

原测试代码（问题版本） #

import pytest
import os
from . import fixtures
from catalog_service import lambda_v1 as v1

@pytest.fixture
def mock_set_environment(monkeypatch):
    monkeypatch.setenv('ENVIRONMENT', 'qa')

class TestLambdaV1:
    @pytest.mark.parametrize('data, expected', [
        (
            {'a': 1},
            {
                'statusCode': '200',
                'headers': {'content-type': 'application/json'},
                'body': '{"a": 1}'
            },
        ),
    ])
    def test_format_good_response(self, mock_set_environment, data, expected):
        assert os.getenv('ENVIRONMENT') == 'qa'
        result = v1.format_good_response(data)
        assert result == expected

核心思路是把模块导入的时机推迟到mock完成之后，同时mock掉boto3和AppSearch的调用，避免真实请求。

方法1：动态导入+局部mock（快速解决） #

在测试函数内部先完成所有mock操作，再导入Lambda模块，确保初始化时用的是模拟好的变量和依赖：

import pytest
import os
from unittest.mock import Mock, patch

class TestLambdaV1:
    @pytest.mark.parametrize('data, expected', [
        (
            {'a': 1},
            {
                'statusCode': '200',
                'headers': {'content-type': 'application/json'},
                'body': '{"a": 1}'
            },
        ),
    ])
    def test_format_good_response(self, monkeypatch):
        # 1. 先模拟环境变量
        monkeypatch.setenv('ENVIRONMENT', 'qa')
        monkeypatch.setenv('APP_SEARCH_API_HOST', 'https://example.aws-appsearch.com')

        # 2. Mock boto3的SSM客户端及get_parameter方法
        mock_ssm_client = Mock()
        mock_ssm_client.get_parameter.return_value = {
            'Parameter': {'Value': 'mock-api-key-123'}
        }
        with patch('boto3.client', return_value=mock_ssm_client):
            # 3. 此时才导入Lambda模块，确保初始化逻辑用的是mock值
            from catalog_service import lambda_v1 as v1

            # 4. 执行测试断言
            assert os.getenv('ENVIRONMENT') == 'qa'
            result = v1.format_good_response(data)
            assert result == expected

            # 验证SSM调用是否符合预期
            mock_ssm_client.get_parameter.assert_called_once_with(
                Name='/qa/app-search/api-key',
                WithDecryption=True
            )

方法2：全局前置mock（多测试用例复用） #

如果多个测试用例都需要相同的前置mock，可以在conftest.py中提前处理，让所有测试执行前就完成环境变量和依赖的模拟：

conftest.py #

import pytest
import os
from unittest.mock import Mock, patch

@pytest.fixture(scope="session", autouse=True)
def mock_lambda_dependencies():
    # 提前设置环境变量
    os.environ['ENVIRONMENT'] = 'qa'
    os.environ['APP_SEARCH_API_HOST'] = 'https://example.aws-appsearch.com'

    # Mock boto3的SSM客户端
    mock_ssm_client = Mock()
    mock_ssm_client.get_parameter.return_value = {
        'Parameter': {'Value': 'mock-api-key-123'}
    }
    # Mock AppSearch类，避免真实初始化请求
    mock_app_search = Mock()
    with patch('boto3.client', return_value=mock_ssm_client), \
         patch('elastic_enterprise_search.AppSearch', return_value=mock_app_search):
        # 提前导入模块，后续测试可直接使用
        import catalog_service.lambda_v1
        yield

简化后的测试代码 #

import pytest
from catalog_service import lambda_v1 as v1

class TestLambdaV1:
    @pytest.mark.parametrize('data, expected', [
        (
            {'a': 1},
            {
                'statusCode': '200',
                'headers': {'content-type': 'application/json'},
                'body': '{"a": 1}'
            },
        ),
    ])
    def test_format_good_response(self, data, expected):
        result = v1.format_good_response(data)
        assert result == expected

方法3：重构Lambda代码（长期最优方案） #

从根源解决问题，把模块级的初始化逻辑移到延迟加载的函数或Lambda handler内部，避免导入时就执行依赖初始化：

重构后的lambda_v1.py #

import boto3
import botocore
import json
import os

from aws_lambda_powertools import Logger
from elastic_enterprise_search import AppSearch


LOGGER = Logger()
# 初始化变量改为None，延迟加载
ENV = None
ssm = None
APP_SEARCH = None

def get_env():
    global ENV
    if not ENV:
        ENV = os.environ['ENVIRONMENT']
    return ENV

def get_app_search():
    global APP_SEARCH, ssm
    if not APP_SEARCH:
        env = get_env()
        ssm = boto3.client('ssm')
        key = f"/{env}/app-search/api-key"
        key_response = ssm.get_parameter(
            Name=key,
            WithDecryption=True
        )
        APP_SEARCH = AppSearch(
            os.environ['APP_SEARCH_API_HOST'],
            bearer_auth=key_response['Parameter']['Value'],
        )
    return APP_SEARCH

# Lambda handler
def lambda_handler(event, context):
    app_search = get_app_search()
    # 业务逻辑...

# 待测试函数
def format_good_response(data):
    return {
        'statusCode': '200',
        'headers': {'content-type': 'application/json'},
        'body': json.dumps(data)
    }

重构后的测试代码 #

这种情况下，初始化逻辑只有在调用get_app_search()或handler时才会执行，测试时可以自由mock：

import pytest
import os
from unittest.mock import Mock, patch
from catalog_service import lambda_v1 as v1

@pytest.fixture
def mock_environment(monkeypatch):
    monkeypatch.setenv('ENVIRONMENT', 'qa')
    monkeypatch.setenv('APP_SEARCH_API_HOST', 'https://example.aws-appsearch.com')

class TestLambdaV1:
    def test_format_good_response(self, mock_environment):
        result = v1.format_good_response({'a':1})
        assert result == {
            'statusCode': '200',
            'headers': {'content-type': 'application/json'},
            'body': '{"a": 1}'
        }

    def test_get_app_search(self, mock_environment):
        # Mock boto3 SSM调用
        mock_ssm_client = Mock()
        mock_ssm_client.get_parameter.return_value = {
            'Parameter': {'Value': 'mock-key'}
        }
        with patch('boto3.client', return_value=mock_ssm_client), \
             patch('elastic_enterprise_search.AppSearch') as mock_app_search_cls:
            app_search = v1.get_app_search()
            # 验证调用逻辑
            mock_ssm_client.get_parameter.assert_called_once_with(
                Name='/qa/app-search/api-key',
                WithDecryption=True
            )
            mock_app_search_cls.assert_called_once_with(
                'https://example.aws-appsearch.com',
                bearer_auth='mock-key'
            )

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