refactor and new test
This commit is contained in:
parent
214835c739
commit
9a8c0e80f6
32
test_xvfb.py
32
test_xvfb.py
@ -1,22 +1,25 @@
|
|||||||
#!/usr/bin/env python
|
#!/usr/bin/env python
|
||||||
|
|
||||||
|
|
||||||
from xvfbwrapper import Xvfb
|
|
||||||
|
|
||||||
import os
|
import os
|
||||||
import unittest
|
import unittest
|
||||||
|
|
||||||
|
from xvfbwrapper import Xvfb
|
||||||
|
|
||||||
|
|
||||||
class TestXvfb(unittest.TestCase):
|
class TestXvfb(unittest.TestCase):
|
||||||
|
|
||||||
def setUp(self):
|
def reset_display(self):
|
||||||
os.environ['DISPLAY'] = ':0'
|
os.environ['DISPLAY'] = ':0'
|
||||||
|
|
||||||
|
def setUp(self):
|
||||||
|
self.reset_display()
|
||||||
|
|
||||||
def test_start(self):
|
def test_start(self):
|
||||||
xvfb = Xvfb()
|
xvfb = Xvfb()
|
||||||
self.addCleanup(xvfb.stop)
|
self.addCleanup(xvfb.stop)
|
||||||
xvfb.start()
|
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)
|
self.assertIsNotNone(xvfb.proc)
|
||||||
|
|
||||||
def test_stop(self):
|
def test_stop(self):
|
||||||
@ -28,10 +31,21 @@ class TestXvfb(unittest.TestCase):
|
|||||||
self.assertIsNone(xvfb.proc)
|
self.assertIsNone(xvfb.proc)
|
||||||
self.assertEqual(orig_display, os.environ['DISPLAY'])
|
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):
|
def test_as_context_manager(self):
|
||||||
orig_display = os.environ['DISPLAY']
|
orig_display = os.environ['DISPLAY']
|
||||||
with Xvfb() as xvfb:
|
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.assertIsNotNone(xvfb.proc)
|
||||||
self.assertIsNone(xvfb.proc)
|
self.assertIsNone(xvfb.proc)
|
||||||
self.assertEqual(orig_display, os.environ['DISPLAY'])
|
self.assertEqual(orig_display, os.environ['DISPLAY'])
|
||||||
@ -46,12 +60,14 @@ class TestXvfb(unittest.TestCase):
|
|||||||
self.assertEqual(w, xvfb.width)
|
self.assertEqual(w, xvfb.width)
|
||||||
self.assertEqual(h, xvfb.height)
|
self.assertEqual(h, xvfb.height)
|
||||||
self.assertEqual(depth, xvfb.colordepth)
|
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)
|
self.assertIsNotNone(xvfb.proc)
|
||||||
|
|
||||||
def test_start_with_arbitrary_kwarg(self):
|
def test_start_with_arbitrary_kwarg(self):
|
||||||
xvfb = Xvfb(nolisten='tcp')
|
xvfb = Xvfb(nolisten='tcp')
|
||||||
self.addCleanup(xvfb.stop)
|
self.addCleanup(xvfb.stop)
|
||||||
xvfb.start()
|
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)
|
self.assertIsNotNone(xvfb.proc)
|
||||||
|
@ -10,7 +10,6 @@
|
|||||||
|
|
||||||
import os
|
import os
|
||||||
import fnmatch
|
import fnmatch
|
||||||
import random
|
|
||||||
import subprocess
|
import subprocess
|
||||||
import tempfile
|
import tempfile
|
||||||
import time
|
import time
|
||||||
@ -23,22 +22,20 @@ class Xvfb:
|
|||||||
self.height = height
|
self.height = height
|
||||||
self.colordepth = colordepth
|
self.colordepth = colordepth
|
||||||
|
|
||||||
if not self._xvfb_exists():
|
if not self.xvfb_exists():
|
||||||
msg = 'Can not find Xvfb. Please install it and try again.'
|
msg = 'Can not find Xvfb. Please install it and try again.'
|
||||||
raise EnvironmentError(msg)
|
raise EnvironmentError(msg)
|
||||||
|
|
||||||
self.xvfb_args = [
|
self.extra_xvfb_args = ['-screen', '0', '{}x{}x{}'.format(
|
||||||
'-screen', '0', '%dx%dx%d' %
|
self.width, self.height, self.colordepth)]
|
||||||
(self.width, self.height, self.colordepth)
|
|
||||||
]
|
|
||||||
|
|
||||||
for key, value in kwargs.items():
|
for key, value in kwargs.items():
|
||||||
self.xvfb_args += ['-%s' % key, value]
|
self.extra_xvfb_args += ['-{}'.format(key), value]
|
||||||
|
|
||||||
if 'DISPLAY' in os.environ:
|
if 'DISPLAY' in os.environ:
|
||||||
self.old_display_num = os.environ['DISPLAY'].split(':')[1]
|
self.orig_display = os.environ['DISPLAY'].split(':')[1]
|
||||||
else:
|
else:
|
||||||
self.old_display_num = 0
|
self.orig_display = None
|
||||||
|
|
||||||
self.proc = None
|
self.proc = None
|
||||||
|
|
||||||
@ -50,25 +47,26 @@ class Xvfb:
|
|||||||
self.stop()
|
self.stop()
|
||||||
|
|
||||||
def start(self):
|
def start(self):
|
||||||
self.vdisplay_num = self._get_next_unused_display()
|
self.new_display = self._get_next_unused_display()
|
||||||
self.xvfb_cmd = ['Xvfb', ':%d' % self.vdisplay_num] + self.xvfb_args
|
display_var = ':{}'.format(self.new_display)
|
||||||
|
self.xvfb_cmd = ['Xvfb', display_var] + self.extra_xvfb_args
|
||||||
with open(os.devnull, 'w') as fnull:
|
with open(os.devnull, 'w') as fnull:
|
||||||
self.proc = subprocess.Popen(self.xvfb_cmd,
|
self.proc = subprocess.Popen(self.xvfb_cmd,
|
||||||
stdout=fnull,
|
stdout=fnull,
|
||||||
stderr=fnull,
|
stderr=fnull,
|
||||||
close_fds=True)
|
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()
|
ret_code = self.proc.poll()
|
||||||
if ret_code is None:
|
if ret_code is None:
|
||||||
self._redirect_display(self.vdisplay_num)
|
self._set_display_var(self.new_display)
|
||||||
else:
|
else:
|
||||||
self._redirect_display(self.old_display_num)
|
|
||||||
self.proc = None
|
|
||||||
raise RuntimeError('Xvfb did not start')
|
raise RuntimeError('Xvfb did not start')
|
||||||
|
|
||||||
def stop(self):
|
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:
|
# TODO:
|
||||||
# fix leaking X displays.
|
# fix leaking X displays.
|
||||||
# (killing Xvfb process doesn't clean up the underlying X11 server)
|
# (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
|
highest_display = max(existing_displays) if existing_displays else 0
|
||||||
return highest_display + 1
|
return highest_display + 1
|
||||||
|
|
||||||
def _redirect_display(self, display_num):
|
def xvfb_exists(self):
|
||||||
os.environ['DISPLAY'] = ':%s' % display_num
|
"""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):
|
def _set_display_var(self, display):
|
||||||
"""Check that Xvfb is in PATH and is executable."""
|
os.environ['DISPLAY'] = ':{}'.format(display)
|
||||||
if any(os.access(os.path.join(path, 'Xvfb'), os.X_OK)
|
|
||||||
for path in os.environ['PATH'].split(os.pathsep))
|
|
||||||
|
Loading…
Reference in New Issue
Block a user