Use a customized virtual display in headless mode

requirements.txt

@@ -29,4 +29,3 @@ pyotp>=2.3.0
 boto>=2.49.0
 flake8>=3.7.8
 certifi>=2019.6.16
-PyVirtualDisplay==0.2.1

seleniumbase/common/unobfuscate.py

@@ -9,15 +9,16 @@
 """
 
 from seleniumbase.common import encryption
+import sys
 import time
 
 
 def main():
-    try:
+    if sys.version_info[0] >= 3:
+        input_method = input  # Using Python 3 (or higher)
+    else:
         # Python 2 has the raw_input() method. Python 3 does not.
         input_method = raw_input  # noqa: ignore=F821
-    except Exception:
-        input_method = input  # Using Python 3
     try:
         while(1):
             code = input_method(

seleniumbase/core/log_helper.py

@@ -33,7 +33,7 @@ def log_test_failure_data(test, test_logpath, driver, browser):
     data_to_save = []
     data_to_save.append("Last_Page: %s" % last_page)
     data_to_save.append("Browser: %s " % browser)
-    if sys.version.startswith('3') and hasattr(test, '_outcome'):
+    if sys.version_info[0] >= 3 and hasattr(test, '_outcome'):
         if test._outcome.errors:
             try:
                 exc_message = test._outcome.errors[0][1][1]

seleniumbase/fixtures/base_case.py

@@ -3256,7 +3256,8 @@ def setUp(self, masterqa_mode=False):
                 self.case_start_time = int(time.time() * 1000)
             if self.headless:
                 try:
-                    from pyvirtualdisplay import Display
+                    # from pyvirtualdisplay import Display  # Skip for own lib
+                    from seleniumbase.virtual_display.display import Display
                     self.display = Display(visible=0, size=(1440, 1880))
                     self.display.start()
                     self.headless_active = True
@@ -3376,7 +3377,7 @@ def tearDown(self):
         super(SubClassOfBaseCase, self).tearDown()
         """
         has_exception = False
-        if sys.version.startswith('3') and hasattr(self, '_outcome'):
+        if sys.version_info[0] >= 3 and hasattr(self, '_outcome'):
             if hasattr(self._outcome, 'errors') and self._outcome.errors:
                 has_exception = True
         else:

seleniumbase/plugins/selenium_plugin.py

@@ -285,7 +285,8 @@ def beforeTest(self, test):
             test.test.headed = True
         if self.options.headless:
             try:
-                from pyvirtualdisplay import Display
+                # from pyvirtualdisplay import Display  # Skip for own lib
+                from seleniumbase.virtual_display.display import Display
                 self.display = Display(visible=0, size=(1440, 1880))
                 self.display.start()
                 self.headless_active = True

seleniumbase/virtual_display/abstractdisplay.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/display.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