Fix a bug that prevented some Xvfb processes from ending

Author: mdmintz

seleniumbase/fixtures/base_case.py

@@ -173,6 +173,7 @@ def __initialize_variables(self):
         self._chart_first_series = {}
         self._chart_series_count = {}
         self._tour_steps = {}
+        self._xvfb_display = None
 
     @classmethod
     def main(self, name, file, *args):
@@ -12721,8 +12722,9 @@ def __activate_virtual_display_as_needed(self):
             try:
                 from sbvirtualdisplay import Display
 
-                self.display = Display(visible=0, size=(width, height))
-                self.display.start()
+                self._xvfb_display = Display(visible=0, size=(width, height))
+                self._xvfb_display.start()
+                sb_config._virtual_display = self._xvfb_display
                 self.headless_active = True
                 sb_config.headless_active = True
             except Exception:
@@ -14833,15 +14835,6 @@ def tearDown(self):
                     self.__activate_debug_mode_in_teardown()
                 # (Pytest) Finally close all open browser windows
                 self.__quit_all_drivers()
-            if self.headless or self.headless2 or self.xvfb:
-                if self.headless_active:
-                    try:
-                        self.display.stop()
-                    except AttributeError:
-                        pass
-                    except Exception:
-                        pass
-                    self.display = None
             if self.with_db_reporting:
                 if has_exception:
                     self.__insert_test_result(constants.State.FAILED, True)
@@ -14924,15 +14917,6 @@ def tearDown(self):
                     print(msg)
                 if self.dashboard:
                     self.__process_dashboard(has_exception)
-                if self.headless or self.headless2 or self.xvfb:
-                    if self.headless_active:
-                        try:
-                            self.display.stop()
-                        except AttributeError:
-                            pass
-                        except Exception:
-                            pass
-                        self.display = None
             if has_exception:
                 test_id = self.__get_test_id()
                 test_logpath = os.path.join(self.log_path, test_id)
@@ -14987,6 +14971,16 @@ def tearDown(self):
             # (Nosetests / Behave / Pure Python) Close all open browser windows
             self.__quit_all_drivers()
         # Resume tearDown() for all test runners, (Pytest / Nosetests / Behave)
+        if hasattr(self, "_xvfb_display") and self._xvfb_display:
+            try:
+                if hasattr(self._xvfb_display, "stop"):
+                    self._xvfb_display.stop()
+                self._xvfb_display = None
+                self.headless_active = False
+            except AttributeError:
+                pass
+            except Exception:
+                pass
         if self.__visual_baseline_copies:
             sb_config._visual_baseline_copies = True
             if has_exception:

seleniumbase/plugins/pytest_plugin.py

@@ -1825,24 +1825,22 @@ def pytest_runtest_teardown(item):
                 except Exception:
                     pass
         try:
-            if hasattr(self, "xvfb") and self.xvfb:
-                if self.headless_active and "--pdb" not in sys_argv:
-                    if hasattr(self, "display") and self.display:
-                        self.headless_active = False
-                        sb_config.headless_active = False
-                        self.display.stop()
-            elif hasattr(self, "headless") and self.headless:
-                if self.headless_active and "--pdb" not in sys_argv:
-                    if hasattr(self, "display") and self.display:
-                        self.headless_active = False
-                        sb_config.headless_active = False
-                        self.display.stop()
-            elif hasattr(self, "headless2") and self.headless2:
-                if self.headless_active and "--pdb" not in sys_argv:
-                    if hasattr(self, "display") and self.display:
-                        self.headless_active = False
-                        sb_config.headless_active = False
-                        self.display.stop()
+            if (
+                hasattr(self, "_xvfb_display")
+                and self._xvfb_display
+                and hasattr(self._xvfb_display, "stop")
+            ):
+                self.headless_active = False
+                sb_config.headless_active = False
+                self._xvfb_display.stop()
+                self._xvfb_display = None
+            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
         except Exception:
             pass
     except Exception:

seleniumbase/plugins/selenium_plugin.py

@@ -1123,9 +1123,6 @@ def beforeTest(self, test):
             # (Set --server="127.0.0.1" for localhost Grid)
             if str(self.options.port) == "443":
                 test.test.protocol = "https"
-        if self.options.xvfb and "linux" not in sys.platform:
-            # The Xvfb virtual display server is for Linux OS Only!
-            self.options.xvfb = False
         if (
             "linux" in sys.platform
             and not self.options.headed
@@ -1171,6 +1168,9 @@ def beforeTest(self, test):
         if not self.options.headless and not self.options.headless2:
             self.options.headed = True
             test.test.headed = True
+        sb_config._virtual_display = None
+        sb_config.headless_active = False
+        self.headless_active = False
         if (
             self.options.headless
             or self.options.headless2
@@ -1180,8 +1180,9 @@ def beforeTest(self, test):
                 # from pyvirtualdisplay import Display  # Skip for own lib
                 from sbvirtualdisplay import Display
 
-                self.display = Display(visible=0, size=(1440, 1880))
-                self.display.start()
+                self._xvfb_display = Display(visible=0, size=(1440, 1880))
+                self._xvfb_display.start()
+                sb_config._virtual_display = self._xvfb_display
                 self.headless_active = True
                 sb_config.headless_active = True
             except Exception:
@@ -1213,13 +1214,22 @@ def afterTest(self, test):
             pass
         except Exception:
             pass
-        if self.options.headless or self.options.xvfb:
-            if self.headless_active:
-                try:
-                    self.headless_active = False
-                    sb_config.headless_active = False
-                    self.display.stop()
-                except AttributeError:
-                    pass
-                except Exception:
-                    pass
+        try:
+            if (
+                hasattr(self, "_xvfb_display")
+                and self._xvfb_display
+                and hasattr(self._xvfb_display, "stop")
+            ):
+                self.headless_active = False
+                sb_config.headless_active = False
+                self._xvfb_display.stop()
+                self._xvfb_display = None
+            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
+        except Exception:
+            pass