Improve support for Xquartz and multi-threaded execution

README.rst

@@ -176,3 +176,32 @@ This test class uses *selenium webdriver* and *xvfbwrapper* to run test cases on
 *Look Ma', no browser!*
 
 (You can also take screenshots inside the virtual display to help diagnose test failures)
+
+----
+
+*******************************************************
+    Example of multi-threaded execution
+*******************************************************
+
+To run several xvfb servers at the same time, you can use the environ keyword
+when starting the Xvfb instances. This provides isolation between threads. Be
+sure to use the environment dictionary you initialize Xvfb with in your
+subsequent system calls. Also, if you wish to inherit your current environment
+you must use the copy method of os.environ and not simply assign a new
+variable to os.environ:
+
+.. code:: python
+
+    from xvfbwrapper import Xvfb
+    import subprocess as sp
+    import os
+
+    isolated_environment = os.environ.copy()
+    xvfb = Xvfb(environ=isolated_environment)
+    xvfb.start()
+    sp.run(
+        "xterm & sleep 1; kill %1 ",
+        shell=True,
+        env=isolated_environment,
+    )
+    xvfb.stop()

test_xvfb.py

@@ -42,6 +42,22 @@ def test_stop(self):
         self.assertEqual(orig_display, os.environ['DISPLAY'])
         self.assertIsNone(xvfb.proc)
 
+    def test_stop_with_xquartz(self):
+        # check that xquartz pattern for display server is dealt with by
+        # xvfb.stop() and restored appropriately
+        xquartz_display = '/private/tmp/com.apple.launchd.CgDzCWvNb1/org.macosforge.xquartz:0'
+        with patch.dict('os.environ',
+           {
+               'DISPLAY':xquartz_display
+           }) as mocked_env:
+
+            xvfb = Xvfb()
+            xvfb.start()
+            self.assertNotEqual(xquartz_display, os.environ['DISPLAY'])
+            xvfb.stop()
+            self.assertEqual(xquartz_display, os.environ['DISPLAY'])
+            self.assertIsNone(xvfb.proc)
+
     def test_start_without_existing_display(self):
         del os.environ['DISPLAY']
         xvfb = Xvfb()
@@ -121,3 +137,23 @@ def test_get_next_unused_display_does_not_reuse_lock(self):
                 self.assertEqual(mockrandint.call_count, 3)
                 self.assertEqual(xvfb3._get_next_unused_display(), 33)
                 self.assertEqual(mockrandint.call_count, 10)
+
+
+    def test_environ_keyword_isolates_environment_modification(self):
+        with patch.dict('os.environ',
+           {
+               'DISPLAY':':0'
+           }) as mocked_env:
+            # Check that start and stop methods modified the environ dict if
+            # passed and does not modify os.environ
+            env_duped = os.environ.copy()
+            xvfb = Xvfb(environ=env_duped)
+            xvfb.start()
+            new_display = ":{}".format(xvfb.new_display)
+            self.assertEqual(':0', os.environ['DISPLAY'])
+            self.assertEqual(new_display, env_duped['DISPLAY'])
+            xvfb.stop()
+            self.assertEqual(':0', os.environ['DISPLAY'])
+            self.assertEqual(':0', env_duped['DISPLAY'])
+            self.assertIsNone(xvfb.proc)
+

xvfbwrapper.py

@@ -32,14 +32,19 @@ class Xvfb(object):
     MAX_DISPLAY = 2147483647
     SLEEP_TIME_BEFORE_START = 0.1
 
-    def __init__(self, width=800, height=680, colordepth=24, tempdir=None, display=None,
+    def __init__(self, width=800, height=680, colordepth=24, tempdir=None, display=None,environ=None,
                  **kwargs):
         self.width = width
         self.height = height
         self.colordepth = colordepth
         self._tempdir = tempdir or tempfile.gettempdir()
         self.new_display = display
 
+        if environ:
+            self.environ = environ
+        else:
+            self.environ = os.environ
+
         if not self.xvfb_exists():
             msg = 'Can not find Xvfb. Please install it and try again.'
             raise EnvironmentError(msg)
@@ -50,10 +55,10 @@ def __init__(self, width=800, height=680, colordepth=24, tempdir=None, display=N
         for key, value in kwargs.items():
             self.extra_xvfb_args += ['-{}'.format(key), value]
 
-        if 'DISPLAY' in os.environ:
-            self.orig_display = os.environ['DISPLAY'].split(':')[1]
+        if 'DISPLAY' in self.environ:
+            self.orig_display_var = self.environ['DISPLAY']
         else:
-            self.orig_display = None
+            self.orig_display_var = None
 
         self.proc = None
 
@@ -81,18 +86,18 @@ def start(self):
         time.sleep(self.__class__.SLEEP_TIME_BEFORE_START)
         ret_code = self.proc.poll()
         if ret_code is None:
-            self._set_display_var(self.new_display)
+            self._set_display(display_var)
         else:
             self._cleanup_lock_file()
             raise RuntimeError('Xvfb did not start ({0}): {1}'
                                .format(ret_code, self.xvfb_cmd))
 
     def stop(self):
         try:
-            if self.orig_display is None:
-                del os.environ['DISPLAY']
+            if self.orig_display_var is None:
+                del self.environ['DISPLAY']
             else:
-                self._set_display_var(self.orig_display)
+                self._set_display(self.orig_display_var)
             if self.proc is not None:
                 try:
                     self.proc.terminate()
@@ -105,7 +110,7 @@ def stop(self):
 
     def xvfb_exists(self):
         """Check that Xvfb is available on PATH and is executable."""
-        paths = os.environ['PATH'].split(os.pathsep)
+        paths = self.environ['PATH'].split(os.pathsep)
         return any(os.access(os.path.join(path, 'Xvfb'), os.X_OK)
                    for path in paths)
 
@@ -161,5 +166,5 @@ def _get_next_unused_display(self):
             else:
                 continue
 
-    def _set_display_var(self, display):
-        os.environ['DISPLAY'] = ':{}'.format(display)
+    def _set_display(self, display_var):
+        self.environ['DISPLAY'] = display_var