如何封装unittest的assertRaises到自定义上下文管理器

原写法失效原因 #

你最初的写法仅通过yield返回了assertRaises生成的上下文管理器对象，但并未激活该对象的拦截逻辑：with wrapper() 语句只会进入wrapper自身的上下文，yield出的对象没有被with语句接管，assertRaises的异常捕获逻辑完全不会触发，抛出的ValueError会直接透传导致程序报错。

可用的基础封装实现 #

要让assertRaises正常捕获异常，必须在自定义上下文管理器内部主动激活assertRaises的上下文，将用户执行的代码块包裹在assertRaises的作用范围内：

import unittest
from contextlib import contextmanager


@contextmanager
def wrapper(expected_exception):
    print("Before assertion")
    # 主动进入assertRaises上下文，将yield放在该上下文内部
    with unittest.TestCase().assertRaises(expected_exception):
        yield
    # 预期异常被assertRaises正常捕获后，会执行后续逻辑
    print("After assertion")


# 测试用例
with wrapper(ValueError):
    raise ValueError()

运行上述代码输出如下，异常被正常捕获，前后自定义逻辑均能正常执行：

Before assertion
After assertion

捕获未被assertRaises拦截的异常 #

你提到的加try/except的思路可行，只需将异常捕获逻辑包裹住assertRaises上下文与yield语句即可，能覆盖两类异常场景：

• 代码块未抛出任何预期异常，此时assertRaises会主动抛出AssertionError
• 代码块抛出了不属于预期类型的其他异常，该异常会从assertRaises透传出来

具体实现代码：

@contextmanager
def wrapper(expected_exception):
    print("Before assertion")
    try:
        with unittest.TestCase().assertRaises(expected_exception):
            yield
        print("After assertion: expected exception caught successfully")
    except Exception as e:
        print(f"Did not raise expected {expected_exception.__name__}! Caught unexpected: {type(e).__name__}: {e}")

对应不同场景的测试效果：

• 场景1：正常抛出预期异常ValueError

  with wrapper(ValueError):
      raise ValueError("test error")
  

  输出：

  Before assertion
  After assertion: expected exception caught successfully
  

• 场景2：代码块未抛出任何异常

  with wrapper(ValueError):
      pass
  

  输出：

  Before assertion
  Did not raise expected ValueError! Caught unexpected: AssertionError: ValueError not raised
  

• 场景3：抛出非预期类型异常

  with wrapper(ValueError):
      raise TypeError("type mismatch")
  

  输出：

  Before assertion
  Did not raise expected ValueError! Caught unexpected: TypeError: type mismatch
  

内容的提问来源于stack exchange，提问作者Mr-Pepe