tests.py

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):
                raise exc
        except Exception as e:
            self.assertEqual(exc, e)
        self.assertEqual(self._history, ["start", "finish"])


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