Code cleanup for test.py.

test.py

@@ -133,27 +133,57 @@ def read_text_file(filename):
     else:
         return open(filename,"r").read()
 
-def random_string(n):
+def wait_until_file_exists(file_path, times=None, interval=0.01):
-    return ''.join(random.choice(string.ascii_lowercase) for x in range(n))
+    while times is None or times:
+        if os.path.exists(file_path):
+            return True
+        time.sleep(interval)
+        if times is not None:
+            times -= 1
+    return False
+
+class TempFileManager(object):
+    def __init__(self, name=""):
+        self._temp_dir = tempfile.mkdtemp(prefix="UltiSnips_" + name)
+
+    def name_temp(self, file_path):
+        return os.path.join(self._temp_dir, file_path)
+
+    def write_temp(self, file_path, content):
+        abs_path = self.name_temp(file_path)
+        create_directory(os.path.dirname(abs_path))
+        if PYTHON3:
+            with open(abs_path, "w", encoding="utf-8") as f:
+                f.write(content)
+        else:
+            with open(abs_path, "w") as f:
+                f.write(content)
+        return abs_path
+
+    def unique_name_temp(self, suffix="", prefix=""):
+        file_handler, abspath = tempfile.mkstemp(suffix, prefix, self._temp_dir)
+        os.close(file_handler)
+        os.remove(abspath)
+        return abspath
+
+    def clear_temp(self):
+        shutil.rmtree(self._temp_dir)
+        create_directory(self._temp_dir)
+
+
+class VimInterface(TempFileManager):
+    def __init__(self, name=""):
+        TempFileManager.__init__(self, name)
 
-class VimInterface(object):
     def get_buffer_data(self):
-        handle, fn = tempfile.mkstemp(prefix="UltiSnips_Test",suffix=".txt")
+        buffer_path = self.unique_name_temp(prefix="buffer_")
-        os.close(handle)
+        self.send(ESC + ":w! %s\n" % buffer_path)
-        os.unlink(fn)
+        if wait_until_file_exists(buffer_path, 50):
-
+            return read_text_file(buffer_path)[:-1]
-        self.send(ESC + ":w! %s\n" % fn)
-
-        # Read the output, chop the trailing newline
-        tries = 50
-        while tries:
-            if os.path.exists(fn):
-                return read_text_file(fn)[:-1]
-            time.sleep(.01)
-            tries -= 1
 
 class VimInterfaceScreen(VimInterface):
     def __init__(self, session):
+        VimInterface.__init__(self, "Screen")
         self.session = session
         self.need_screen_escapes = 0
         self.detect_parsing()
@@ -184,15 +214,14 @@ class VimInterfaceScreen(VimInterface):
         # Send a string where the interpretation will depend on version of screen
         string = "$TERM"
         self.send("i" + string + ESC)
-        # too fast makes the buffer differ from string even in
-        # non-escape mode
-        time.sleep(1)
         output = self.get_buffer_data()
         # If the output doesn't match the input, need to do additional escaping
         if output != string:
             self.need_screen_escapes = 1
         self.send(ESC + ":q!\n")
-        time.sleep(0.2) # prevent from losing the first vim test case
+        # Make sure that vim is closed, preventing from losing the first
+        # vim test case
+        time.sleep(0.2)
 
 class VimInterfaceTmux(VimInterface):
     def __init__(self, session):
@@ -284,20 +313,8 @@ class VimInterfaceWindows(VimInterface):
 
         self.shell.SendKeys(seq)
 
-def create_temp_file(prefix, suffix, content):
-    """Create a file in a temporary place with the given 'prefix'
-    and the given 'suffix' containing 'content'. The file is never
-    deleted. Returns the name of the temporary file."""
-    with tempfile.NamedTemporaryFile(
-        prefix=prefix, suffix=suffix, delete=False
-    ) as temporary_file:
-        if PYTHON3:
-            s = s.encode("utf-8")
-        temporary_file.write(content)
-        temporary_file.close()
-        return temporary_file.name
 
-class _VimTest(unittest.TestCase):
+class _VimTest(unittest.TestCase, TempFileManager):
     snippets = ()
     files = {}
     text_before = " --- some text before --- \n\n"
@@ -312,6 +329,10 @@ class _VimTest(unittest.TestCase):
     skip_if = lambda self: None
     version = None  # Will be set to vim --version output
 
+    def __init__(self, *args, **kwargs):
+        unittest.TestCase.__init__(self, *args, **kwargs)
+        TempFileManager.__init__(self, "Test")
+
     def runTest(self):
         # Only checks the output. All work is done in setUp().
         wanted = self.text_before + self.wanted + self.text_after
@@ -340,21 +361,11 @@ class _VimTest(unittest.TestCase):
     def _create_file(self, file_path, content):
         """Creates a file in the runtimepath that is created for this test.
         Returns the absolute path to the file."""
-        abs_path = os.path.join(self._temporary_directory, *file_path.split("/"))
+        return self.write_temp(file_path, dedent(content + "\n"))
-        create_directory(os.path.dirname(abs_path))
-
-        content = dedent(content + "\n")
-        if PYTHON3:
-            with open(abs_path, "w", encoding="utf-8") as file_handle:
-                file_handle.write(content)
-        else:
-            with open(abs_path, "w") as file_handle:
-                file_handle.write(content)
-        return abs_path
 
     def _link_file(self, source, relative_destination):
         """Creates a link from 'source' to the 'relative_destination' in our temp dir."""
-        absdir = os.path.join(self._temporary_directory, relative_destination)
+        absdir = self.name_temp(relative_destination)
         create_directory(absdir)
         os.symlink(source, os.path.join(absdir, os.path.basename(source)))
 
@@ -371,11 +382,9 @@ class _VimTest(unittest.TestCase):
         if reason_for_skipping is not None:
             return self.skipTest(reason_for_skipping)
 
-        self._temporary_directory = tempfile.mkdtemp(prefix="UltiSnips_Test")
-
         vim_config = []
         vim_config.append('set nocompatible')
-        vim_config.append('set runtimepath=$VIMRUNTIME,.,%s' % self._temporary_directory)
+        vim_config.append('set runtimepath=$VIMRUNTIME,.,%s' % self._temp_dir)
 
         if self.plugins:
             self._link_file(os.path.join(plugin_cache_dir(), "vim-pathogen", "autoload"), ".")
@@ -430,8 +439,8 @@ class _VimTest(unittest.TestCase):
         vim_config.append("vim.current.window.cursor = (max(len(vim.current.buffer)//2, 1), 0)")
 
         # Create a file to signalize to the test runner that we are done with starting Vim.
-        vim_pid_file = os.path.join(self._temporary_directory, "vim.pid")
+        vim_pid_file = self.name_temp("vim.pid")
-        done_file = os.path.join(self._temporary_directory, "loading_done")
+        done_file = self.name_temp("loading_done")
         vim_config.append("with open('%s', 'w') as pid_file: pid_file.write(vim.eval('getpid()'))" %
                 vim_pid_file)
         vim_config.append("with open('%s', 'w') as done_file: pass" % done_file)
@@ -443,15 +452,11 @@ class _VimTest(unittest.TestCase):
             self._create_file(name, content)
 
         # Now launch Vim.
-        self._create_file("vim_config.vim", os.linesep.join(vim_config))
+        vim_config_path = self._create_file("vim_config.vim", os.linesep.join(vim_config))
         # Note the shell to exclude it from shell history.
-        self.vim.send(""" vim -u %s\r\n""" % os.path.join(
+        self.vim.send(""" vim -u %s\r\n""" % vim_config_path)
-            self._temporary_directory, "vim_config.vim"))
+        wait_until_file_exists(done_file)
-        while True:
-            if os.path.exists(done_file):
         self._vim_pid = int(open(vim_pid_file, "r").read())
-                break
-            time.sleep(.01)
 
         self._before_test()
 
@@ -465,12 +470,12 @@ class _VimTest(unittest.TestCase):
 
     def tearDown(self):
         if self.interrupt:
-            print("Working directory: %s" % (self._temporary_directory))
+            print("Working directory: %s" % (self._temp_dir))
             return
-        shutil.rmtree(self._temporary_directory)
         self.vim.send(3*ESC + ":qa!\n")
         while is_process_running(self._vim_pid):
             time.sleep(.05)
+        self.clear_temp()
 
 ###########################################################################
 #                            BEGINNING OF TEST                            #
@@ -3385,8 +3390,7 @@ class MySnippetSource(SnippetSource):
     return []
 """)
         pyfile = 'py3file' if PYTHON3 else 'pyfile'
-        vim_config.append("%s %s" % (pyfile, os.path.join(
+        vim_config.append("%s %s" % (pyfile, self.name_temp("snippet_source.py")))
-            self._temporary_directory, "snippet_source.py")))
 # End: Snippet Source  #}}}
 
 # Plugin: YouCompleteMe  {{{#