Add virtual_display package to SeleniumBase

Author: mdmintz

seleniumbase/virtual_display/__init__.py

@@ -0,0 +1,139 @@
+import fnmatch
+import os
+import tempfile
+import time
+from threading import Lock
+from seleniumbase.virtual_display.easyprocess import EasyProcess
+from seleniumbase.virtual_display import xauth
+
+mutex = Lock()
+RANDOMIZE_DISPLAY_NR = False
+if RANDOMIZE_DISPLAY_NR:
+    import random
+    random.seed()
+MIN_DISPLAY_NR = 1000
+USED_DISPLAY_NR_LIST = []
+
+
+class AbstractDisplay(EasyProcess):
+    '''
+    Common parent for Xvfb and Xephyr
+    '''
+    def __init__(self, use_xauth=False):
+        mutex.acquire()
+        try:
+            self.display = self.search_for_display()
+            while self.display in USED_DISPLAY_NR_LIST:
+                self.display += 1
+            USED_DISPLAY_NR_LIST.append(self.display)
+        finally:
+            mutex.release()
+        if xauth and not xauth.is_installed():
+            raise xauth.NotFoundError()
+        self.use_xauth = use_xauth
+        self._old_xauth = None
+        self._xauth_filename = None
+        EasyProcess.__init__(self, self._cmd)
+
+    @property
+    def new_display_var(self):
+        return ':%s' % (self.display)
+
+    @property
+    def _cmd(self):
+        raise NotImplementedError()
+
+    def lock_files(self):
+        tmpdir = '/tmp'
+        pattern = '.X*-lock'
+        # remove path.py dependency
+        names = fnmatch.filter(os.listdir(tmpdir), pattern)
+        ls = [os.path.join(tmpdir, child) for child in names]
+        ls = [p for p in ls if os.path.isfile(p)]
+        return ls
+
+    def search_for_display(self):
+        # search for free display
+        ls = [int(x.split('X')[1].split('-')[0]) for x in self.lock_files()]
+        if len(ls):
+            display = max(MIN_DISPLAY_NR, max(ls) + 3)
+        else:
+            display = MIN_DISPLAY_NR
+        if RANDOMIZE_DISPLAY_NR:
+            display += random.randint(0, 100)
+        return display
+
+    def redirect_display(self, on):
+        '''
+        on:
+         * True -> set $DISPLAY to virtual screen
+         * False -> set $DISPLAY to original screen
+
+        :param on: bool
+        '''
+        d = self.new_display_var if on else self.old_display_var
+        if d is None:
+            del os.environ['DISPLAY']
+        else:
+            os.environ['DISPLAY'] = d
+
+    def start(self):
+        '''
+        start display
+
+        :rtype: self
+        '''
+        if self.use_xauth:
+            self._setup_xauth()
+        EasyProcess.start(self)
+
+        # https://github.com/ponty/PyVirtualDisplay/issues/2
+        # https://github.com/ponty/PyVirtualDisplay/issues/14
+        self.old_display_var = os.environ.get('DISPLAY', None)
+
+        self.redirect_display(True)
+        # wait until X server is active
+        # TODO: better method
+        time.sleep(0.1)
+        return self
+
+    def stop(self):
+        '''
+        stop display
+
+        :rtype: self
+        '''
+        self.redirect_display(False)
+        EasyProcess.stop(self)
+        if self.use_xauth:
+            self._clear_xauth()
+        return self
+
+    def _setup_xauth(self):
+        '''
+        Set up the Xauthority file and the XAUTHORITY environment variable.
+        '''
+        handle, filename = tempfile.mkstemp(prefix='PyVirtualDisplay.',
+                                            suffix='.Xauthority')
+        self._xauth_filename = filename
+        os.close(handle)
+        # Save old environment
+        self._old_xauth = {}
+        self._old_xauth['AUTHFILE'] = os.getenv('AUTHFILE')
+        self._old_xauth['XAUTHORITY'] = os.getenv('XAUTHORITY')
+
+        os.environ['AUTHFILE'] = os.environ['XAUTHORITY'] = filename
+        cookie = xauth.generate_mcookie()
+        xauth.call('add', self.new_display_var, '.', cookie)
+
+    def _clear_xauth(self):
+        '''
+        Clear the Xauthority file and restore the environment variables.
+        '''
+        os.remove(self._xauth_filename)
+        for varname in ['AUTHFILE', 'XAUTHORITY']:
+            if self._old_xauth[varname] is None:
+                del os.environ[varname]
+            else:
+                os.environ[varname] = self._old_xauth[varname]
+        self._old_xauth = None

seleniumbase/virtual_display/abstractdisplay.py

@@ -0,0 +1,61 @@
+"""
+This module contains a customized version of pyvirtualdisplay.
+These helper methods SHOULD NOT be called directly from tests.
+"""
+from seleniumbase.virtual_display.abstractdisplay import AbstractDisplay
+from seleniumbase.virtual_display.xephyr import XephyrDisplay
+from seleniumbase.virtual_display.xvfb import XvfbDisplay
+from seleniumbase.virtual_display.xvnc import XvncDisplay
+
+
+class Display(AbstractDisplay):
+    '''
+    Common class
+
+    :param color_depth: [8, 16, 24, 32]
+    :param size: screen size (width,height)
+    :param bgcolor: background color ['black' or 'white']
+    :param visible: True -> Xephyr, False -> Xvfb
+    :param backend: 'xvfb', 'xvnc' or 'xephyr', ignores ``visible``
+    :param xauth: If a Xauthority file should be created.
+    '''
+    def __init__(self, backend=None, visible=False, size=(1024, 768),
+                 color_depth=24, bgcolor='black', use_xauth=False, **kwargs):
+        self.color_depth = color_depth
+        self.size = size
+        self.bgcolor = bgcolor
+        self.screen = 0
+        self.process = None
+        self.display = None
+        self.visible = visible
+        self.backend = backend
+
+        if not self.backend:
+            if self.visible:
+                self.backend = 'xephyr'
+            else:
+                self.backend = 'xvfb'
+
+        self._obj = self.display_class(
+            size=size,
+            color_depth=color_depth,
+            bgcolor=bgcolor,
+            **kwargs)
+        AbstractDisplay.__init__(self, use_xauth=use_xauth)
+
+    @property
+    def display_class(self):
+        assert self.backend
+        if self.backend == 'xvfb':
+            cls = XvfbDisplay
+        if self.backend == 'xvnc':
+            cls = XvncDisplay
+        if self.backend == 'xephyr':
+            cls = XephyrDisplay
+        cls.check_installed()
+        return cls
+
+    @property
+    def _cmd(self):
+        self._obj.display = self.display
+        return self._obj._cmd