python/tests: add TestCaseInTempdir.mktemp()

This gives you a name of a temporary file within the test case's tempdir.
Use it like this:

  with self.mktemp() as filename:
     self.check_run("samba-tool foo --output %s" % filename)
     self.assertStringsEqual(open(filename).read(), expected)

and filename will flick out of existence when the with block ends.

This is based on an idea used in the traffic_runner tests, which will
soon be adapted to use this method.

Signed-off-by: Douglas Bagnall <douglas.bagnall@catalyst.net.nz>
Reviewed-by: Andrew Bartlett <abartlet@samba.org>

diff --git a/python/samba/tests/__init__.py b/python/samba/tests/__init__.py
index e95b6a9157af8bdc38f68c4ba5dea583d2a5229b..f904499b90b7c7c7ccc4f9d88bd29e099fd00df4 100644 (file)
--- a/python/samba/tests/__init__.py
+++ b/python/samba/tests/__init__.py
@@ -17,7 +17,7 @@
 #
 
 """Samba Python tests."""
-
+from __future__ import print_function
 import os
 import tempfile
 import warnings
@@ -37,6 +37,7 @@ from samba.compat import text_type
 from samba.compat import string_types
 from random import randint
 from random import SystemRandom
+from contextlib import contextmanager
 import string
 try:
     from samba.samdb import SamDB
@@ -299,6 +300,20 @@ class TestCaseInTempDir(TestCase):
         os.rmdir(self.tempdir)
         self.tempdir = None
 
+    @contextmanager
+    def mktemp(self):
+        """Yield a temporary filename in the tempdir."""
+        try:
+            fd, fn = tempfile.mkstemp(dir=self.tempdir)
+            yield fn
+        finally:
+            try:
+                os.close(fd)
+                os.unlink(fn)
+            except (OSError, IOError) as e:
+                print("could not remove temporary file: %s" % e,
+                      file=sys.stderr)
+
 
 def env_loadparm():
     lp = param.LoadParm()