refactor and new test

test_xvfb.py

@@ -1,22 +1,25 @@
 #!/usr/bin/env python
 
-
-from xvfbwrapper import Xvfb
-
 import os
 import unittest
 
+from xvfbwrapper import Xvfb
+
 
 class TestXvfb(unittest.TestCase):
 
-    def setUp(self):
+    def reset_display(self):
         os.environ['DISPLAY'] = ':0'
 
+    def setUp(self):
+        self.reset_display()
+
     def test_start(self):
         xvfb = Xvfb()
         self.addCleanup(xvfb.stop)
         xvfb.start()
-        self.assertEqual(':%d' % xvfb.vdisplay_num, os.environ['DISPLAY'])
+        display_var = ':{}'.format(xvfb.new_display)
+        self.assertEqual(display_var, os.environ['DISPLAY'])
         self.assertIsNotNone(xvfb.proc)
 
     def test_stop(self):
@@ -28,10 +31,21 @@ class TestXvfb(unittest.TestCase):
         self.assertIsNone(xvfb.proc)
         self.assertEqual(orig_display, os.environ['DISPLAY'])
 
+    def test_start_without_existing_display(self):
+        del os.environ['DISPLAY']
+        xvfb = Xvfb()
+        self.addCleanup(xvfb.stop)
+        self.addCleanup(self.reset_display)
+        xvfb.start()
+        display_var = ':{}'.format(xvfb.new_display)
+        self.assertEqual(display_var, os.environ['DISPLAY'])
+        self.assertIsNotNone(xvfb.proc)
+
     def test_as_context_manager(self):
         orig_display = os.environ['DISPLAY']
         with Xvfb() as xvfb:
-            self.assertEqual(':%d' % xvfb.vdisplay_num, os.environ['DISPLAY'])
+            display_var = ':{}'.format(xvfb.new_display)
+            self.assertEqual(display_var, os.environ['DISPLAY'])
             self.assertIsNotNone(xvfb.proc)
         self.assertIsNone(xvfb.proc)
         self.assertEqual(orig_display, os.environ['DISPLAY'])
@@ -46,12 +60,14 @@ class TestXvfb(unittest.TestCase):
         self.assertEqual(w, xvfb.width)
         self.assertEqual(h, xvfb.height)
         self.assertEqual(depth, xvfb.colordepth)
-        self.assertEqual(os.environ['DISPLAY'], ':%d' % xvfb.vdisplay_num)
+        display_var = ':{}'.format(xvfb.new_display)
+        self.assertEqual(display_var, os.environ['DISPLAY'])
         self.assertIsNotNone(xvfb.proc)
 
     def test_start_with_arbitrary_kwarg(self):
         xvfb = Xvfb(nolisten='tcp')
         self.addCleanup(xvfb.stop)
         xvfb.start()
-        self.assertEqual(os.environ['DISPLAY'], ':%d' % xvfb.vdisplay_num)
+        display_var = ':{}'.format(xvfb.new_display)
+        self.assertEqual(display_var, os.environ['DISPLAY'])
         self.assertIsNotNone(xvfb.proc)

xvfbwrapper.py

@@ -10,7 +10,6 @@
 
 import os
 import fnmatch
-import random
 import subprocess
 import tempfile
 import time
@@ -23,22 +22,20 @@ class Xvfb:
         self.height = height
         self.colordepth = colordepth
 
-        if not self._xvfb_exists():
+        if not self.xvfb_exists():
             msg = 'Can not find Xvfb. Please install it and try again.'
             raise EnvironmentError(msg)
 
-        self.xvfb_args = [
-            '-screen', '0', '%dx%dx%d' %
-            (self.width, self.height, self.colordepth)
-        ]
+        self.extra_xvfb_args = ['-screen', '0', '{}x{}x{}'.format(
+                                self.width, self.height, self.colordepth)]
 
         for key, value in kwargs.items():
-            self.xvfb_args += ['-%s' % key, value]
+            self.extra_xvfb_args += ['-{}'.format(key), value]
 
         if 'DISPLAY' in os.environ:
-            self.old_display_num = os.environ['DISPLAY'].split(':')[1]
+            self.orig_display = os.environ['DISPLAY'].split(':')[1]
         else:
-            self.old_display_num = 0
+            self.orig_display = None
 
         self.proc = None
 
@@ -50,25 +47,26 @@ class Xvfb:
         self.stop()
 
     def start(self):
-        self.vdisplay_num = self._get_next_unused_display()
-        self.xvfb_cmd = ['Xvfb', ':%d' % self.vdisplay_num] + self.xvfb_args
-
+        self.new_display = self._get_next_unused_display()
+        display_var = ':{}'.format(self.new_display)
+        self.xvfb_cmd = ['Xvfb', display_var] + self.extra_xvfb_args
         with open(os.devnull, 'w') as fnull:
             self.proc = subprocess.Popen(self.xvfb_cmd,
                                          stdout=fnull,
                                          stderr=fnull,
                                          close_fds=True)
-        time.sleep(0.2)  # give Xvfb time to start
+        time.sleep(0.1)  # give Xvfb time to start
         ret_code = self.proc.poll()
         if ret_code is None:
-            self._redirect_display(self.vdisplay_num)
+            self._set_display_var(self.new_display)
         else:
-            self._redirect_display(self.old_display_num)
-            self.proc = None
             raise RuntimeError('Xvfb did not start')
 
     def stop(self):
-        self._redirect_display(self.old_display_num)
+        if self.orig_display is None:
+            del os.environ['DISPLAY']
+        else:
+            self._set_display_var(self.orig_display)
         # TODO:
         # fix leaking X displays.
         # (killing Xvfb process doesn't clean up the underlying X11 server)
@@ -86,10 +84,11 @@ class Xvfb:
         highest_display = max(existing_displays) if existing_displays else 0
         return highest_display + 1
 
-    def _redirect_display(self, display_num):
-        os.environ['DISPLAY'] = ':%s' % display_num
+    def xvfb_exists(self):
+        """Check that Xvfb is available on PATH and is executable."""
+        paths = os.environ['PATH'].split(os.pathsep)
+        return any(os.access(os.path.join(path, 'Xvfb'), os.X_OK)
+                   for path in paths)
 
-    def _xvfb_exists(self):
-        """Check that Xvfb is in PATH and is executable."""
-        if any(os.access(os.path.join(path, 'Xvfb'), os.X_OK)
-               for path in os.environ['PATH'].split(os.pathsep))
+    def _set_display_var(self, display):
+        os.environ['DISPLAY'] = ':{}'.format(display)