diff --git a/Lib/test/test_contextlib.py b/Lib/test/test_contextlib.py index b878dc1614f..c05f37b0ba4 100644 --- a/Lib/test/test_contextlib.py +++ b/Lib/test/test_contextlib.py @@ -100,6 +100,128 @@ class ContextManagerTestCase(unittest.TestCase): self.assertEqual(baz.foo, 'bar') self.assertEqual(baz.__doc__, "Whee!") +class NestedTestCase(unittest.TestCase): + + # XXX This needs more work + + def test_nested(self): + @contextmanager + def a(): + yield 1 + @contextmanager + def b(): + yield 2 + @contextmanager + def c(): + yield 3 + with nested(a(), b(), c()) as (x, y, z): + self.assertEqual(x, 1) + self.assertEqual(y, 2) + self.assertEqual(z, 3) + + def test_nested_cleanup(self): + state = [] + @contextmanager + def a(): + state.append(1) + try: + yield 2 + finally: + state.append(3) + @contextmanager + def b(): + state.append(4) + try: + yield 5 + finally: + state.append(6) + try: + with nested(a(), b()) as (x, y): + state.append(x) + state.append(y) + 1/0 + except ZeroDivisionError: + self.assertEqual(state, [1, 4, 2, 5, 6, 3]) + else: + self.fail("Didn't raise ZeroDivisionError") + + def test_nested_right_exception(self): + state = [] + @contextmanager + def a(): + yield 1 + class b(object): + def __enter__(self): + return 2 + def __exit__(self, *exc_info): + try: + raise Exception() + except: + pass + try: + with nested(a(), b()) as (x, y): + 1/0 + except ZeroDivisionError: + self.assertEqual((x, y), (1, 2)) + except Exception: + self.fail("Reraised wrong exception") + else: + self.fail("Didn't raise ZeroDivisionError") + + def test_nested_b_swallows(self): + @contextmanager + def a(): + yield + @contextmanager + def b(): + try: + yield + except: + # Swallow the exception + pass + try: + with nested(a(), b()): + 1/0 + except ZeroDivisionError: + self.fail("Didn't swallow ZeroDivisionError") + + def test_nested_break(self): + @contextmanager + def a(): + yield + state = 0 + while True: + state += 1 + with nested(a(), a()): + break + state += 10 + self.assertEqual(state, 1) + + def test_nested_continue(self): + @contextmanager + def a(): + yield + state = 0 + while state < 3: + state += 1 + with nested(a(), a()): + continue + state += 10 + self.assertEqual(state, 3) + + def test_nested_return(self): + @contextmanager + def a(): + try: + yield + except: + pass + def foo(): + with nested(a(), a()): + return 1 + return 10 + self.assertEqual(foo(), 1) + class ClosingTestCase(unittest.TestCase): # XXX This needs more work