如何为含os.popen环境变量的main.py编写单元测试并Mock环境

原代码文件：main.py #

import os
_ENV = os.popen("echo ${RPM_ENVIRONMENT}").read().split('\n')[0]
_HOST = os.popen("echo $(hostname)").read().split('\n')[0]
_JOB_TYPE=os.popen("echo ${JOB_TYPE}").read().split('\n')[0]

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

问题分析 #

导入main.py时触发KeyError: '${RPM_ENVIRONMENT}'，原因是系统未配置RPM_ENVIRONMENT环境变量，导致_ENV取值为字符串${RPM_ENVIRONMENT}，无法匹配SERVER_URL字典的键。需要在导入前Mockos.popen的返回值，让_ENV被设为DEV。

修正后的测试文件：test_main.py #

import unittest
import sys
from unittest.mock import patch

sys.path.insert(1, 'C:/home/src')

def mock_popen(cmd):
    # 根据执行的命令返回对应模拟结果
    if "${RPM_ENVIRONMENT}" in cmd:
        # 返回DEV，模拟环境变量值
        return type('MockPopen', (), {'read': lambda: 'DEV\n'})()
    elif "$(hostname)" in cmd:
        # 模拟主机名返回
        return type('MockPopen', (), {'read': lambda: 'test-host\n'})()
    elif "${JOB_TYPE}" in cmd:
        # 模拟JOB_TYPE返回
        return type('MockPopen', (), {'read': lambda: 'TEST\n'})()

# 必须在导入main前完成Mock，因为导入时会立即执行环境变量读取逻辑
with patch('os.popen', side_effect=mock_popen):
    import main as conf

class TestMainConfig(unittest.TestCase):
    def test_server_url_correct(self):
        # 验证SERVER_URL返回DEV环境的地址
        self.assertEqual(conf.SERVER_URL.get('ENV_URL'), 'https://dev.net')
    
    def test_env_value(self):
        # 验证_ENV变量被正确设为DEV
        self.assertEqual(conf._ENV, 'DEV')

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

核心要点 #

• 导入前Mock：main.py在导入阶段就会执行os.popen读取环境变量，所以必须在import main之前完成Mock，否则会先触发KeyError
• 命令匹配模拟：通过side_effect根据不同的命令字符串返回对应结果，确保每个os.popen调用都能得到合理的模拟值
• 测试用例覆盖：分别验证目标变量SERVER_URL和_ENV的取值符合预期

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