From 9a8c0e80f6a932d5448d1b97d5a9c316a7a7c8e3 Mon Sep 17 00:00:00 2001 From: Corey Goldberg Date: Sat, 2 Jan 2016 21:20:24 -0500 Subject: [PATCH] refactor and new test --- test_xvfb.py | 32 ++++++++++++++++++++++++-------- xvfbwrapper.py | 45 ++++++++++++++++++++++----------------------- 2 files changed, 46 insertions(+), 31 deletions(-) diff --git a/test_xvfb.py b/test_xvfb.py index 8aa0530..13d0bfe 100644 --- a/test_xvfb.py +++ b/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) diff --git a/xvfbwrapper.py b/xvfbwrapper.py index 5612fb1..11739e1 100644 --- a/xvfbwrapper.py +++ b/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)