diff --git a/src/context_manager.py b/src/context_manager.py new file mode 100644 index 0000000..27a7a43 --- /dev/null +++ b/src/context_manager.py @@ -0,0 +1,17 @@ +from contextlib import contextmanager + + +class ClassyContextManager: + def __init__(self, thing): + self.thing = thing + + def __enter__(self): + return self.thing + + def __exit__(self, exc_type, exc_val, exc_tb): + pass + + +@contextmanager +def funky_context_manager(thing): + yield thing diff --git a/tests/test_context_manager.py b/tests/test_context_manager.py new file mode 100644 index 0000000..2e68ae0 --- /dev/null +++ b/tests/test_context_manager.py @@ -0,0 +1,15 @@ +from unittest.mock import Mock + +from context_manager import ClassyContextManager, funky_context_manager + + +def test_class_context_manager(): + mock = Mock() + with ClassyContextManager(mock) as classy: + assert classy is mock + + +def test_funky_context_manager(): + mock = Mock() + with funky_context_manager(mock) as funky: + assert funky is mock