Browse Source

refactor and new test

master
Corey Goldberg 6 years ago
parent
commit
9a8c0e80f6
  1. 32
      test_xvfb.py
  2. 45
      xvfbwrapper.py

32
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)

45
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)
Loading…
Cancel
Save