diff --git a/Lib/slapdtest/_slapdtest.py b/Lib/slapdtest/_slapdtest.py index f248ff68..f1885caf 100644 --- a/Lib/slapdtest/_slapdtest.py +++ b/Lib/slapdtest/_slapdtest.py @@ -164,6 +164,14 @@ class SlapdObject(object): When a reference to an instance of this class is lost, the slapd server is shut down. + + An instance can be used as a context manager. When exiting the context + manager, the slapd server is shut down and the temporary data store is + removed. + + .. versionchanged:: 3.1 + + Added context manager functionality """ slapd_conf_template = SLAPD_CONF_TEMPLATE database = 'mdb' @@ -553,6 +561,13 @@ def ldapdelete(self, dn, recursive=False, extra_args=None): extra_args.append(dn) self._cli_popen(self.PATH_LDAPDELETE, extra_args=extra_args) + def __enter__(self): + self.start() + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.stop() + class SlapdTestCase(unittest.TestCase): """ diff --git a/Tests/t_slapdobject.py b/Tests/t_slapdobject.py new file mode 100644 index 00000000..e1cc971b --- /dev/null +++ b/Tests/t_slapdobject.py @@ -0,0 +1,18 @@ +import unittest + +import slapdtest + + +class TestSlapdObject(unittest.TestCase): + def test_context_manager(self): + with slapdtest.SlapdObject() as server: + self.assertIsNotNone(server._proc) + self.assertIsNone(server._proc) + + def test_context_manager_after_start(self): + server = slapdtest.SlapdObject() + server.start() + self.assertIsNotNone(server._proc) + with server: + self.assertIsNotNone(server._proc) + self.assertIsNone(server._proc)