import asyncio
import unittest


class TestAsyncIteratorFinally(unittest.TestCase):
    # finally: blocks in __aiter__ should execute
    # even if exceptions are raised.

    class AsyncFinallyIterator:

        def __init__(self):
            self.finally_executed = False

        async def __aiter__(self):
            for i in range(10):
                try:
                    yield i
                finally:
                    self.finally_executed = True

    async def main(self, it):
        async for i in it:
            raise RuntimeError()

    def test_main(self):
        it = self.AsyncFinallyIterator()
        loop = asyncio.get_event_loop()
        with self.assertRaises(RuntimeError):
            loop.run_until_complete(self.main(it))
        self.assertTrue(it.finally_executed)


class TestIteratorFinally(unittest.TestCase):
    # finally: blocks in __iter__ should execute
    # even if exceptions are raised.

    class FinallyIterator:

        def __init__(self):
            self.finally_executed = False

        def __iter__(self):
            for i in range(10):
                try:
                    yield i
                finally:
                    self.finally_executed = True

    def main(self, it):
        for _ in it:
            raise RuntimeError()

    def test_main(self):
        it = self.FinallyIterator()
        with self.assertRaises(RuntimeError):
            self.main(it)
        self.assertTrue(it.finally_executed)


if __name__ == '__main__':
    unittest.main()