如何在pytest中设置环境变量解决KeyError: 'BUCKET'报错问题

报错的根本原因是业务代码里的DATA_ENGINEERING_BUCKET = os.environ["BUCKET"]是模块级代码，会在你导入runner.kinesis模块的时候直接执行。而你现有测试写法里的monkeypatch.setenv是在测试方法运行阶段才执行，此时模块导入已经完成，环境变量还没设置，所以直接抛出KeyError，根本不会走到测试方法的设置逻辑。
另外你当前的写法还有一个问题：monkeypatch是pytest的fixture，不能直接作为参数传入unittest.TestCase子类的测试方法中，运行时会报参数不匹配的错误。

方案1：导入业务模块前预先设置环境变量 #

直接在测试文件的最顶部、导入runner.kinesis之前就设置好环境变量，保证模块导入时能读取到值：

import os
# 先设置环境变量再导入业务模块
os.environ["BUCKET"] = "test-bucket"
from runner import kinesis
import unittest

class TestHandler(unittest.TestCase):
    @mock_s3
    @mock_lambda
    def test_handler(self):
        actual = kinesis.handler(kinesis_stream_event, "")
        expected = {"statusCode": 200, "body": "OK"}
        assert actual == expected

如果要全局生效，可以在pytest的conftest.py中添加会话级自动运行的fixture提前设置：

import pytest
import os

@pytest.fixture(scope="session", autouse=True)
def preset_env():
    os.environ["BUCKET"] = "test-bucket"

方案2：修改业务代码为惰性加载环境变量 #

把模块级的环境变量读取逻辑延后到handler调用时再执行，避免导入阶段就读取：

# 业务代码修改后
import os

def handler(event, context):
    DATA_ENGINEERING_BUCKET = os.environ["BUCKET"]
    # 原有逻辑不变

如果多个函数需要用到该变量，可以封装成工具函数实现懒加载：

import os

def get_data_bucket():
    return os.environ["BUCKET"]

def handler(event, context):
    DATA_ENGINEERING_BUCKET = get_data_bucket()
    # 原有逻辑不变

修改后你原有的测试逻辑只需要调整环境变量的设置方式即可：

class TestHandler(unittest.TestCase):
    @mock_s3
    @mock_lambda
    def test_handler(self):
        os.environ["BUCKET"] = "test-bucket"
        actual = kinesis.handler(kinesis_stream_event, "")
        expected = {"statusCode": 200, "body": "OK"}
        assert actual == expected
        # 测试完成后清理环境变量
        del os.environ["BUCKET"]

方案3：测试时重新加载模块 #

如果你不想修改业务代码，也可以在测试方法中先设置环境变量，再重新加载业务模块触发模块级代码重新执行：

import importlib
import os
import unittest

class TestHandler(unittest.TestCase):
    @mock_s3
    @mock_lambda
    def test_handler(self):
        os.environ["BUCKET"] = "test-bucket"
        # 重新加载模块，此时环境变量已经设置完成
        import runner.kinesis
        importlib.reload(runner.kinesis)
        actual = runner.kinesis.handler(kinesis_stream_event, "")
        expected = {"statusCode": 200, "body": "OK"}
        assert actual == expected
        del os.environ["BUCKET"]

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