unittest2.case.assertRaises 的實作方式

New and Improved Coming changes to unittest in Python 2.7 & 3.2 有詳細的 unittest2 介紹, 其中我覺得最漂亮的改進是 assertRaises, 用法變成:

with self.assertRaises(TypeError) as cm:
    do_something()

exception = cm.exception
self.assertEqual(exception.error_code, 3)

我覺得新介面用 with 執行, 既易懂又能彈性地執行不同行數的程式。想不到它怎麼做出這樣的行為。好奇之下, 看了一下原始碼, 才發覺作法很簡單:

def assertRaises(self, excClass, callableObj=None, *args, **kwargs):
        # ...
        if callableObj is None:
            return _AssertRaisesContext(excClass, self)
        # ...

class _AssertRaisesContext(object):
    """A context manager used to implement TestCase.assertRaises* methods."""

    def __init__(self, expected, test_case, expected_regexp=None):
        self.expected = expected
        self.failureException = test_case.failureException
        self.expected_regexp = expected_regexp

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, tb):
        if exc_type is None:
            try:
                exc_name = self.expected.__name__
            except AttributeError:
                exc_name = str(self.expected)
            raise self.failureException(
                "%s not raised" % (exc_name,))
        # ...

就是傳回一個物件 _AssertRaisesContext, 讓該物件實作 with 的 __enter__ 和 __exit__, __exit__ 的參數含有過程中產生的 exception 相關資訊, 離開時檢查它就知道有沒有產生 exception 了。

另外 unittest2 為了和 unittest 相容, 介面寫得挺漂亮的, unittest2 的 assertRaises 會先依參數數量判斷是舊的呼叫方式還是新的使用 with 的呼叫方式, 再做對應的處理。需要維護的程式變多後, 對於向下相容的藝術, 才有進一步的體會。