使用MagicMock测试含上下文管理器的Python类时的属性异常问题

待测试代码 #

from typing import Callable, Any

class HelperClass:
    @property
    def output(self) -> dict[str, Any]:
        return self._output

    def __enter__(self):
        print("__enter__")
        self._output: dict[str, Any] = {}
        return self

    def __exit__(self, exc_type, exc_value, exc_traceback):
        print("__exit__")


class ImportantClass:
    def fn(self, key: str):
        with self._fac() as hc:
            if key not in hc.output:
                raise RuntimeError("bad bad not good")

            return hc.output[key]

    def __init__(self, fac: Callable[[], HelperClass]):
        self._fac = fac

测试代码及问题 #

原测试代码尝试用MagicMock模拟HelperClass，并通过PropertyMock设置output属性，但在ImportantClass.fn方法中，hc.output变成了MagicMock实例而非预期的字典，导致断言失败：

from unittest.mock import MagicMock, PropertyMock

def test_important_class():
    hc = MagicMock(spec=HelperClass)
    omock = PropertyMock(return_value={"foo": "bar", "baz": "quux"})
    type(hc).output = omock

    assert hc.output["foo"] == "bar"
    assert hc.output["baz"] == "quux" # 这些断言成功

    assert ImportantClass(fac=lambda: hc).fn("foo") == "bar"
    assert ImportantClass(fac=lambda: hc).fn("baz") == "quux" # 这些断言失败

问题核心：with语句会调用Mock的__enter__方法，默认情况下MagicMock的__enter__会返回一个全新的MagicMock实例，而非我们预先配置了PropertyMock的hc实例，因此fn方法中访问的hc.output是新实例的属性，没有绑定我们定义的PropertyMock。

解决方案 #

修改Mock的__enter__方法，让它返回自身，确保with语句拿到的是同一个已配置好的Mock实例：

from unittest.mock import MagicMock, PropertyMock

def test_important_class():
    hc = MagicMock(spec=HelperClass)
    # 让__enter__返回自身，保证with语句获取的是预先配置的Mock实例
    hc.__enter__.return_value = hc
    
    omock = PropertyMock(return_value={"foo": "bar", "baz": "quux"})
    type(hc).output = omock

    # 验证直接访问逻辑
    assert hc.output["foo"] == "bar"
    assert hc.output["baz"] == "quux"

    # 测试ImportantClass的业务逻辑
    important = ImportantClass(fac=lambda: hc)
    assert important.fn("foo") == "bar"
    assert important.fn("baz") == "quux"

补充说明 #

• 若需要自定义__exit__的行为，可以设置hc.__exit__.return_value = True（默认返回值即为True，通常无需额外配置）
• 使用spec=HelperClass可以约束Mock仅包含HelperClass的属性和方法，避免拼写错误或调用不存在的成员

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