From ba28d2b3bed98d078a0c80345a2f2ed021d90014 Mon Sep 17 00:00:00 2001
From: Paul Sokolovsky <pfalcon@users.sourceforge.net>
Date: Tue, 12 May 2015 13:07:36 +0300
Subject: [PATCH] unittest: Let use assertRaises() as a context manager.

---
 unittest/unittest.py | 21 ++++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

diff --git a/unittest/unittest.py b/unittest/unittest.py
index f8db373f..c3a4b79a 100644
--- a/unittest/unittest.py
+++ b/unittest/unittest.py
@@ -2,6 +2,22 @@ class SkipTest(Exception):
     pass
 
 
+class AssertRaisesContext:
+
+    def __init__(self, exc):
+        self.expected = exc
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, tb):
+        if exc_type is None:
+            assert False, "%r not raised" % exc
+        if issubclass(exc_type, self.expected):
+            return True
+        return False
+
+
 class TestCase:
 
     def fail(self, msg=''):
@@ -26,7 +42,10 @@ class TestCase:
     def assertIsInstance(self, x, y, msg=''):
         assert isinstance(x, y), msg
 
-    def assertRaises(self, exc, func, *args, **kwargs):
+    def assertRaises(self, exc, func=None, *args, **kwargs):
+        if func is None:
+            return AssertRaisesContext(exc)
+
         try:
             func(*args, **kwargs)
             assert False, "%r not raised" % exc