import unittest
from ucontextlib import contextmanager


class ContextManagerTestCase(unittest.TestCase):

    def setUp(self):
        self._history = []

        @contextmanager
        def manager(x):
            self._history.append('start')
            try:
                yield x
            finally:
                self._history.append('finish')

        self._manager = manager

    def test_context_manager(self):
        with self._manager(123) as x:
            self.assertEqual(x, 123)
        self.assertEqual(self._history, ['start', 'finish'])

    def test_context_manager_on_error(self):
        exc = Exception()
        try:
            with self._manager(123) as x:
                raise exc
        except Exception as e:
            self.assertEqual(exc, e)
        self.assertEqual(self._history, ['start', 'finish'])


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