使用Python unittest mock os.environ时的测试失败问题

我正在测试一个处理工作目录的类，它通过类变量map映射目录路径，其中"baz"对应的路径从环境变量BAZ_PATH获取。使用Python 3.8.13和unittest框架测试，希望避免mockWorkingDirectory.map，也不想直接修改os.environ的值。尝试用@mock.patch.dict(os.environ, {"BAZ_PATH": DUMMY_BAZ_PATH})进行测试，但wd.path返回None或实际环境变量值，而非预设的DUMMY_BAZ_PATH，test_controlled_baz测试无法通过。

相关代码 #

working_directory.py #

import os


class WorkingDirectory:
    map = {
        "foo": "path/to/foo",
        "bar": "path/to/bar",
        "baz": os.environ.get("BAZ_PATH"),
    }

    def __init__(self, env: str):
        self.env = env
        self.path = self.map[self.env]

    @property
    def data_dir(self):
        return os.path.join(self.path, "data")

    # Other similar methods...

test.py #

import os
import unittest

from unittest import mock

from working_directory import WorkingDirectory

DUMMY_BAZ_PATH = "path/to/baz"


class TestWorkingDirectory(unittest.TestCase):
    @mock.patch.dict(os.environ, {"BAZ_PATH": DUMMY_BAZ_PATH})
    def test_controlled_baz(self):
        wd = WorkingDirectory("baz")
        self.assertEqual(wd.path, DUMMY_BAZ_PATH)

错误信息 #

======================================================================
FAIL: test_controlled_baz (test_directory_structure_utils.TestWorkingDirectory)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "~/.pyenv/versions/3.8.13/lib/python3.8/unittest/mock.py", line 1756, in _inner
    return f(*args, **kw)
  File "~/Projects/dummy_project/tests/unit/test_directory_structure_utils.py", line 127, in test_controlled_baz
    self.assertEqual(wd.path, DUMMY_BAZ_PATH)
AssertionError: None != 'path/to/baz'

----------------------------------------------------------------------
Ran 136 tests in 0.325s

FAILED (failures=1, skipped=5)

解决方案 #

问题核心是类变量map在模块导入阶段就已初始化，此时os.environ.get("BAZ_PATH")已经执行完成，后续通过mock.patch.dict修改环境变量不会改变已经生成的map值。

方法1：延迟环境变量读取（推荐） #

修改原类，让baz的路径在实例化时动态读取，而不是模块导入时就固定：

# working_directory.py
import os


class WorkingDirectory:
    _static_paths = {
        "foo": "path/to/foo",
        "bar": "path/to/bar",
    }

    def __init__(self, env: str):
        self.env = env
        # 每次实例化时动态构建映射表
        self.path_map = self._static_paths.copy()
        self.path_map["baz"] = os.environ.get("BAZ_PATH")
        self.path = self.path_map[self.env]

    @property
    def data_dir(self):
        return os.path.join(self.path, "data")

这样每次创建WorkingDirectory实例时，都会重新读取环境变量，mock的数值就能生效。

更简洁的写法可以用类方法动态生成映射：

class WorkingDirectory:
    @classmethod
    def _get_path_map(cls):
        return {
            "foo": "path/to/foo",
            "bar": "path/to/bar",
            "baz": os.environ.get("BAZ_PATH"),
        }

    def __init__(self, env: str):
        self.env = env
        self.path = self._get_path_map()[self.env]

方法2：测试时重新导入模块 #

如果不想修改原代码，可以在mock环境变量后，重新导入模块以触发类变量的重新初始化：

# test.py
import os
import unittest
import importlib
from unittest import mock

import working_directory

DUMMY_BAZ_PATH = "path/to/baz"


class TestWorkingDirectory(unittest.TestCase):
    @mock.patch.dict(os.environ, {"BAZ_PATH": DUMMY_BAZ_PATH})
    def test_controlled_baz(self):
        # 重新导入模块，让类变量重新初始化
        importlib.reload(working_directory)
        WorkingDirectory = working_directory.WorkingDirectory
        
        wd = WorkingDirectory("baz")
        self.assertEqual(wd.path, DUMMY_BAZ_PATH)

注意这种方法可能影响其他测试用例的模块状态，需要谨慎使用。

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