Fix --xvfb compatibility with --reuse-session / --rs

Author: mdmintz

seleniumbase/core/browser_launcher.py

@@ -888,6 +888,13 @@ def __install_pyautogui_if_missing():
                     _xvfb_display.start()
                     sb_config._virtual_display = _xvfb_display
                     sb_config.headless_active = True
+                    if (
+                        hasattr(sb_config, "reuse_session")
+                        and sb_config.reuse_session
+                        and hasattr(sb_config, "_vd_list")
+                        and isinstance(sb_config._vd_list, list)
+                    ):
+                        sb_config._vd_list.append(_xvfb_display)
 
 
 def install_pyautogui_if_missing(driver):

seleniumbase/fixtures/base_case.py

@@ -14010,6 +14010,9 @@ def __activate_standard_virtual_display(self):
             if not self.undetectable:
                 sb_config._virtual_display = self._xvfb_display
                 sb_config.headless_active = True
+            if self._reuse_session and hasattr(sb_config, "_vd_list"):
+                if isinstance(sb_config._vd_list, list):
+                    sb_config._vd_list.append(self._xvfb_display)
 
     def __activate_virtual_display(self):
         if self.undetectable and not (self.headless or self.headless2):
@@ -14034,6 +14037,9 @@ def __activate_virtual_display(self):
                     self.__activate_standard_virtual_display()
                 else:
                     self.headless_active = True
+                    if self._reuse_session and hasattr(sb_config, "_vd_list"):
+                        if isinstance(sb_config._vd_list, list):
+                            sb_config._vd_list.append(self._xvfb_display)
             except Exception as e:
                 if hasattr(e, "msg"):
                     print("\n" + str(e.msg))
@@ -14088,7 +14094,7 @@ def __activate_virtual_display_as_needed(self):
         """This is only needed on Linux.
         The "--xvfb" arg is still useful, as it prevents headless mode,
         which is the default mode on Linux unless using another arg."""
-        if "linux" in sys.platform and (not self.headed or self.xvfb):
+        if is_linux and (not self.headed or self.xvfb):
             pip_find_lock = fasteners.InterProcessLock(
                 constants.PipInstall.FINDLOCK
             )
@@ -16605,7 +16611,11 @@ def tearDown(self):
             # (Pynose / Behave / Pure Python) Close all open browser windows
             self.__quit_all_drivers()
         # Resume tearDown() for all test runners, (Pytest / Pynose / Behave)
-        if hasattr(self, "_xvfb_display") and self._xvfb_display:
+        if (
+            hasattr(self, "_xvfb_display")
+            and self._xvfb_display
+            and not self._reuse_session
+        ):
             # Stop the Xvfb virtual display launched from BaseCase
             try:
                 if hasattr(self._xvfb_display, "stop"):
@@ -16620,6 +16630,13 @@ def tearDown(self):
             hasattr(sb_config, "_virtual_display")
             and sb_config._virtual_display
             and hasattr(sb_config._virtual_display, "stop")
+            and (
+                not hasattr(sb_config, "reuse_session")
+                or (
+                    hasattr(sb_config, "reuse_session")
+                    and not sb_config.reuse_session
+                )
+            )
         ):
             # CDP Mode may launch a 2nd Xvfb virtual display
             try:

seleniumbase/plugins/pytest_plugin.py

@@ -1371,6 +1371,7 @@ def pytest_addoption(parser):
 
     arg_join = " ".join(sys_argv)
     sb_config._browser_shortcut = None
+    sb_config._vd_list = []
 
     # SeleniumBase does not support pytest-timeout due to hanging browsers.
     for arg in sys_argv:
@@ -2017,6 +2018,13 @@ def pytest_runtest_teardown(item):
                 hasattr(self, "_xvfb_display")
                 and self._xvfb_display
                 and hasattr(self._xvfb_display, "stop")
+                and (
+                    not hasattr(sb_config, "reuse_session")
+                    or (
+                        hasattr(sb_config, "reuse_session")
+                        and not sb_config.reuse_session
+                    )
+                )
             ):
                 self.headless_active = False
                 sb_config.headless_active = False
@@ -2026,6 +2034,13 @@ def pytest_runtest_teardown(item):
                 hasattr(sb_config, "_virtual_display")
                 and sb_config._virtual_display
                 and hasattr(sb_config._virtual_display, "stop")
+                and (
+                    not hasattr(sb_config, "reuse_session")
+                    or (
+                        hasattr(sb_config, "reuse_session")
+                        and not sb_config.reuse_session
+                    )
+                )
             ):
                 sb_config._virtual_display.stop()
                 sb_config._virtual_display = None
@@ -2139,6 +2154,21 @@ def _perform_pytest_unconfigure_(config):
             except Exception:
                 pass
         sb_config.shared_driver = None
+        with suppress(Exception):
+            if (
+                hasattr(sb_config, "_virtual_display")
+                and sb_config._virtual_display
+                and hasattr(sb_config._virtual_display, "stop")
+            ):
+                sb_config._virtual_display.stop()
+                sb_config._virtual_display = None
+                sb_config.headless_active = False
+            if hasattr(sb_config, "_vd_list") and sb_config._vd_list:
+                if isinstance(sb_config._vd_list, list):
+                    for display in sb_config._vd_list:
+                        if display:
+                            with suppress(Exception):
+                                display.stop()
     if hasattr(sb_config, "log_path") and sb_config.item_count > 0:
         log_helper.archive_logs_if_set(
             constants.Logs.LATEST + "/", sb_config.archive_logs