Repository: bokeh/bokeh

Issue:
Terminate firefox after usage in io.export
I've noticed that there is bunch of orphaned processes left behind on CI.


Discussion:

Resolving patch:
diff --git a/bokeh/io/webdriver.py b/bokeh/io/webdriver.py
--- a/bokeh/io/webdriver.py
+++ b/bokeh/io/webdriver.py
@@ -27,7 +27,7 @@
 import atexit
 import shutil
 from os.path import devnull
-from typing import Optional
+from typing import List, Optional
 
 # External imports
 from selenium import webdriver
@@ -84,39 +84,25 @@ def _try_create_chromium_webdriver() -> Optional[WebDriver]:
     except Exception:
         return None
 
-def _has_firefox() -> bool:
-    return _try_create_firefox_webdriver() is not None
-
-def _has_chromium() -> bool:
-    return _try_create_chromium_webdriver() is not None
-
 class _WebdriverState(object):
     '''
 
     '''
 
-    _reuse: bool
-    _kind: DriverKind
+    reuse: bool
+    kind: Optional[DriverKind]
 
     current: Optional[WebDriver]
+    _drivers: List[WebDriver]
 
-    def __init__(self, *, reuse: bool = True, kind: Optional[DriverKind] = None):
+    def __init__(self, *, kind: Optional[DriverKind] = None, reuse: bool = True):
+        self.kind = kind
         self.reuse = reuse
         self.current = None
+        self._drivers = set()
 
-        if kind is not None:
-            self.kind = kind
-        else:
-            if _has_chromium():
-                self.kind = "chromium"
-            elif _has_firefox():
-                self.kind = "firefox"
-            else:
-                raise RuntimeError("Neither firefox/geckodriver nor chromium-browser/chromedriver are available on system PATH. " \
-                                   "You can install the former with 'conda install -c conda-forge firefox geckodriver'.")
-
-    @staticmethod
-    def terminate(driver: WebDriver) -> None:
+    def terminate(self, driver: WebDriver) -> None:
+        self._drivers.remove(driver)
         driver.quit()
 
     def reset(self) -> None:
@@ -126,35 +112,43 @@ def reset(self) -> None:
 
     def get(self) -> WebDriver:
         if not self.reuse or self.current is None:
-            if self.current is not None:
-                self.terminate(self.current)
+            self.reset()
             self.current = self.create()
         return self.current
 
     def create(self, kind: Optional[DriverKind] = None) -> WebDriver:
+        driver = self._create(kind)
+        self._drivers.add(driver)
+        return driver
+
+    def _create(self, kind: Optional[DriverKind]) -> WebDriver:
         driver_kind = kind or self.kind
-        if driver_kind == "firefox":
-            return create_firefox_webdriver()
-        elif driver_kind == "chromium":
-            return create_chromium_webdriver()
-        else:
-            raise ValueError(f"Unknown webdriver kind {driver_kind}")
 
-    @property
-    def reuse(self) -> bool:
-        return self._reuse
+        if driver_kind is None:
+            driver = _try_create_chromium_webdriver()
+            if driver is not None:
+                self.kind = "chromium"
+                return driver
 
-    @reuse.setter
-    def reuse(self, value: bool) -> None:
-        self._reuse = value
+            driver = _try_create_firefox_webdriver()
+            if driver is not None:
+                self.kind = "firefox"
+                return driver
 
-    @property
-    def kind(self) -> DriverKind:
-        return self._kind
+            raise RuntimeError("Neither firefox and geckodriver nor a variant of chromium browser and " \
+                               "chromedriver are available on system PATH. You can install the former " \
+                               "with 'conda install -c conda-forge firefox geckodriver'.")
+        elif driver_kind == "chromium":
+            return create_chromium_webdriver()
+        elif driver_kind == "firefox":
+            return create_firefox_webdriver()
+        else:
+            raise ValueError(f"'{driver_kind}' is not a recognized webdriver kind")
 
-    @kind.setter
-    def kind(self, value: DriverKind) -> None:
-        self._kind = value
+    def cleanup(self) -> None:
+        self.reset()
+        for driver in list(self._drivers):
+            self.terminate(driver)
 
 #-----------------------------------------------------------------------------
 # Code
@@ -162,4 +156,4 @@ def kind(self, value: DriverKind) -> None:
 
 webdriver_control = _WebdriverState()
 
-atexit.register(lambda: webdriver_control.reset())
+atexit.register(lambda: webdriver_control.cleanup())


Test patch:
diff --git a/tests/unit/bokeh/io/test_webdriver.py b/tests/unit/bokeh/io/test_webdriver.py
--- a/tests/unit/bokeh/io/test_webdriver.py
+++ b/tests/unit/bokeh/io/test_webdriver.py
@@ -35,15 +35,21 @@
 
 def test_create_firefox_webdriver() -> None:
     d = biw.create_firefox_webdriver()
-    assert isinstance(d, selenium.webdriver.firefox.webdriver.WebDriver)
+    try:
+        assert isinstance(d, selenium.webdriver.firefox.webdriver.WebDriver)
+    finally:
+        d.quit()
 
 def test_create_chromium_webdriver() -> None:
     d = biw.create_chromium_webdriver()
-    assert isinstance(d, selenium.webdriver.chrome.webdriver.WebDriver)
+    try:
+        assert isinstance(d, selenium.webdriver.chrome.webdriver.WebDriver)
+    finally:
+        d.quit()
 
 _driver_map = {
-    'firefox': selenium.webdriver.firefox.webdriver.WebDriver,
-    'chromium': selenium.webdriver.chrome.webdriver.WebDriver,
+    "firefox": selenium.webdriver.firefox.webdriver.WebDriver,
+    "chromium": selenium.webdriver.chrome.webdriver.WebDriver,
 }
 
 class Test_webdriver_control(object):
@@ -52,7 +58,7 @@ def test_default(self) -> None:
         # so create a new instance only to check default values
         wc = biw._WebdriverState()
         assert wc.reuse == True
-        assert wc.kind == "chromium"
+        assert wc.kind == None
         assert wc.current is None
 
     def test_get_with_reuse(self) -> None: