* Trying to bring in line a lot more with some of the other Python Projects

[ejpi] / src / libraries / recipes / test_utils.py

diff --git a/src/libraries/recipes/test_utils.py b/src/libraries/recipes/test_utils.py
index 349eded..d238bfa 100644 (file)
--- a/src/libraries/recipes/test_utils.py
+++ b/src/libraries/recipes/test_utils.py
@@ -82,11 +82,12 @@ def CheckReferences(func):
        return wrapper
 
 
-class expected(object):
+@contextlib.contextmanager
+def expected(exception):
        """
-       >>> with expected(ZeroDivisionError):
+       >>> with expected2(ZeroDivisionError):
        ...     1 / 0
-       >>> with expected(AssertionError("expected ZeroDivisionError to have been thrown")):
+       >>> with expected2(AssertionError("expected ZeroDivisionError to have been thrown")):
        ...     with expected(ZeroDivisionError):
        ...             1 / 2
        Traceback (most recent call last):
@@ -97,33 +98,33 @@ class expected(object):
                File "/media/data/Personal/Development/bzr/Recollection-trunk/src/libraries/recipes/context.py", line 139, in __exit__
                        assert t is not None, ("expected {0:%s} to have been thrown" % (self._t.__name__))
        AssertionError: expected {0:ZeroDivisionError} to have been thrown
-       >>> with expected(Exception("foo")):
+       >>> with expected2(Exception("foo")):
        ...     raise Exception("foo")
-       >>> with expected(Exception("bar")):
+       >>> with expected2(Exception("bar")):
        ...     with expected(Exception("foo")): # this won't catch it
        ...             raise Exception("bar")
        ...     assert False, "should not see me"
-       >>> with expected(Exception("can specify")):
+       >>> with expected2(Exception("can specify")):
        ...     raise Exception("can specify prefixes")
-       >>> with expected(Exception("Base class fun")):
-       ...     raise IndexError("Base class fun")
+       >>> with expected2(Exception("Base class fun")):
+       True
+       >>> True
+       False
        """
-
-       def __init__(self, e):
-               if isinstance(e, Exception):
-                       self._t, self._v = type(e), str(e)
-               elif isinstance(e, type):
-                       self._t, self._v = e, ""
-               else:
-                       raise Exception("usage: with expected(Exception): ... or "
-                                                       "with expected(Exception(\"text\")): ...")
-
-       def __enter__(self):
-               try:
-                       pass
-               except:
-                       pass # this is a Python 3000 way of saying sys.exc_clear()
-
-       def __exit__(self, t, v, tb):
-               assert t is not None, ("expected {0:%s} to have been thrown" % (self._t.__name__))
-               return self._t in inspect.getmro(t) and str(v).startswith(self._v)
+       if isinstance(exception, Exception):
+               excType, excValue = type(exception), str(exception)
+       elif isinstance(e, type):
+               excType, excValue = e, ""
+
+       try:
+               yield
+       except Exception, e:
+               if not (excType in inspect.getmro(type(e)) and str(e).startswith(excValue)):
+                       raise
+       else:
+               raise AssertionError("expected {0:%s} to have been thrown" % excType.__name__)
+
+
+if __name__ == "__main__":
+       import doctest
+       doctest.testmod()
\ No newline at end of file