Make SlapdObject a context manager

jdufresne · jdufresne · commit 55c96b995365 · 2018-04-01T14:08:49.000-07:00
Allows for automatic cleanup of resources using the with statement. For
example:

   with slapdtest.SlapdObject() as server:
       server.ldapadd(...)
       ...

When using SlapdObject in a function, it is more convenient and concise
than using try/finally pattern.
diff --git a/Lib/slapdtest/_slapdtest.py b/Lib/slapdtest/_slapdtest.py
@@ -164,6 +164,10 @@ 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.
     """
     slapd_conf_template = SLAPD_CONF_TEMPLATE
     database = 'mdb'
@@ -553,6 +557,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)
