如何在Python单元测试中模拟环境变量RPM_ENVIRONMENT

问题核心在于主文件的顶级代码会在导入时立即执行，_ENV和SERVER_URL的赋值逻辑在import阶段就已运行，因此必须在导入主文件前完成os.popen的模拟，否则会因未正确Mock导致KeyError。

正确测试代码写法 #

在test_env.py中，先完成os.popen的Mock，再导入主模块：

import unittest
from unittest import mock

# 先Mock os.popen，再导入主文件
with mock.patch("os.popen") as popen_mock:
    # 模拟完整调用链：popen().read().split('\n')[0] 返回'DEV'
    # 注意要包含换行符，匹配原代码中split('\n')[0]的逻辑
    popen_mock.return_value.read.return_value = "DEV\n"
    # 此时导入主文件，顶级代码会使用Mock后的os.popen
    import main  # 替换为你的主文件实际文件名

class TestEnv(unittest.TestCase):
    def test_server_url_matches_dev(self):
        self.assertEqual(main.SERVER_URL['ENV_URL'], 'https://dev.net')

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

原写法失效原因 #

如果将Mock逻辑放在测试函数内部，主文件会在文件顶部import时就执行了顶级代码，此时Mock还未生效，_ENV会被赋值为MagicMock对象，最终触发字典查找的KeyError。必须将Mock放在主文件import之前，确保顶级代码执行时Mock已生效。

额外优化建议（可选） #

主文件通过os.popen调用Shell读取环境变量的方式不够高效可靠，建议改用os.getenv直接读取环境变量，测试时会更简单：

修改主文件代码： #

import os

# 直接读取环境变量，可设置默认值避免环境未定义时出错
_ENV = os.getenv('RPM_ENVIRONMENT', 'DEV')

SERVER_URL = {
    'DEV':{'ENV_URL':'https://dev.net'},
    'UAT':{'ENV_URL':'https://uat.net'},
    'PROD':{'ENV_URL':'https://prod.net'}
}[_ENV]

对应测试代码： #

import unittest
import os
from unittest import mock

# 通过Mock环境变量字典，再导入主文件
with mock.patch.dict(os.environ, {'RPM_ENVIRONMENT': 'DEV'}):
    import main

class TestEnv(unittest.TestCase):
    def test_server_url_matches_dev(self):
        self.assertEqual(main.SERVER_URL['ENV_URL'], 'https://dev.net')

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