Merge pull request python-ldap#198 – Make SlapdObject a context manager

encukou · web-flow · commit 2a6bcd41db3f · 2018-04-03T13:39:10.000+02:00
python-ldap#198
diff --git 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
@@ -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)
