[py] Allow driver path to be set using ENV variables (#14528)

Delta456 · harsha509 · web-flow · commit a2cacc197ff0 · 2024-10-04T12:34:34.000+01:00
---------

Co-authored-by: Sri Harsha &lt;12621691+harsha509@users.noreply.github.com&gt;
diff --git a/py/selenium/webdriver/chromium/service.py b/py/selenium/webdriver/chromium/service.py
@@ -39,9 +39,11 @@ def __init__(
         service_args: typing.Optional[typing.List[str]] = None,
         log_output: SubprocessStdAlias = None,
         env: typing.Optional[typing.Mapping[str, str]] = None,
+        driver_path_env_key: str = None,
         **kwargs,
     ) -> None:
         self.service_args = service_args or []
+        driver_path_env_key = driver_path_env_key or "SE_CHROMEDRIVER"
 
         if isinstance(log_output, str):
             self.service_args.append(f"--log-path={log_output}")
@@ -56,6 +58,7 @@ def __init__(
             port=port,
             env=env,
             log_output=self.log_output,
+            driver_path_env_key=driver_path_env_key,
             **kwargs,
         )
 
diff --git a/py/selenium/webdriver/chromium/webdriver.py b/py/selenium/webdriver/chromium/webdriver.py
@@ -51,7 +51,7 @@ def __init__(
             options.binary_location = finder.get_browser_path()
             options.browser_version = None
 
-        self.service.path = finder.get_driver_path()
+        self.service.path = self.service.env_path() or finder.get_driver_path()
         self.service.start()
 
         executor = ChromiumRemoteConnection(
diff --git a/py/selenium/webdriver/common/service.py b/py/selenium/webdriver/common/service.py
@@ -25,6 +25,7 @@
 from platform import system
 from subprocess import PIPE
 from time import sleep
+from typing import Optional
 from typing import cast
 from urllib import request
 from urllib.error import URLError
@@ -53,6 +54,7 @@ def __init__(
         port: int = 0,
         log_output: SubprocessStdAlias = None,
         env: typing.Optional[typing.Mapping[typing.Any, typing.Any]] = None,
+        driver_path_env_key: str = None,
         **kwargs,
     ) -> None:
         if isinstance(log_output, str):
@@ -64,12 +66,13 @@ def __init__(
         else:
             self.log_output = log_output
 
-        self._path = executable_path
         self.port = port or utils.free_port()
         # Default value for every python subprocess: subprocess.Popen(..., creationflags=0)
         self.popen_kw = kwargs.pop("popen_kw", {})
         self.creation_flags = self.popen_kw.pop("creation_flags", 0)
         self.env = env or os.environ
+        self.DRIVER_PATH_ENV_KEY = driver_path_env_key
+        self._path = self.env_path() or executable_path
 
     @property
     def service_url(self) -> str:
@@ -236,3 +239,6 @@ def _start_process(self, path: str) -> None:
                     f"'{os.path.basename(self._path)}' executable may have wrong permissions."
                 ) from err
             raise
+
+    def env_path(self) -> Optional[str]:
+        return os.getenv(self.DRIVER_PATH_ENV_KEY, None)
diff --git a/py/selenium/webdriver/edge/service.py b/py/selenium/webdriver/edge/service.py
@@ -40,15 +40,18 @@ def __init__(
         log_output: SubprocessStdAlias = None,
         service_args: typing.Optional[typing.List[str]] = None,
         env: typing.Optional[typing.Mapping[str, str]] = None,
+        driver_path_env_key: str = None,
         **kwargs,
     ) -> None:
         self.service_args = service_args or []
+        driver_path_env_key = driver_path_env_key or "SE_EDGEDRIVER"
 
         super().__init__(
             executable_path=executable_path,
             port=port,
             service_args=service_args,
             log_output=log_output,
             env=env,
+            driver_path_env_key=driver_path_env_key,
             **kwargs,
         )
diff --git a/py/selenium/webdriver/firefox/service.py b/py/selenium/webdriver/firefox/service.py
@@ -40,15 +40,18 @@ def __init__(
         service_args: typing.Optional[typing.List[str]] = None,
         log_output: SubprocessStdAlias = None,
         env: typing.Optional[typing.Mapping[str, str]] = None,
+        driver_path_env_key: str = None,
         **kwargs,
     ) -> None:
         self.service_args = service_args or []
+        driver_path_env_key = driver_path_env_key or "SE_GECKODRIVER"
 
         super().__init__(
             executable_path=executable_path,
             port=port,
             log_output=log_output,
             env=env,
+            driver_path_env_key=driver_path_env_key,
             **kwargs,
         )
 
diff --git a/py/selenium/webdriver/firefox/webdriver.py b/py/selenium/webdriver/firefox/webdriver.py
@@ -58,7 +58,7 @@ def __init__(
             options.binary_location = finder.get_browser_path()
             options.browser_version = None
 
-        self.service.path = finder.get_driver_path()
+        self.service.path = self.service.env_path() or finder.get_driver_path()
         self.service.start()
 
         executor = FirefoxRemoteConnection(
diff --git a/py/selenium/webdriver/ie/service.py b/py/selenium/webdriver/ie/service.py
@@ -32,6 +32,7 @@ def __init__(
         service_args: typing.Optional[typing.List[str]] = None,
         log_level: typing.Optional[str] = None,
         log_output: SubprocessStdAlias = None,
+        driver_path_env_key: str = None,
         **kwargs,
     ) -> None:
         """Creates a new instance of the Service.
@@ -46,6 +47,8 @@ def __init__(
            Default is "stdout".
         """
         self.service_args = service_args or []
+        driver_path_env_key = driver_path_env_key or "SE_IEDRIVER"
+
         if host:
             self.service_args.append(f"--host={host}")
         if log_level:
@@ -55,6 +58,7 @@ def __init__(
             executable_path=executable_path,
             port=port,
             log_output=log_output,
+            driver_path_env_key=driver_path_env_key,
             **kwargs,
         )
 
diff --git a/py/selenium/webdriver/ie/webdriver.py b/py/selenium/webdriver/ie/webdriver.py
@@ -46,7 +46,7 @@ def __init__(
         self.service = service if service else Service()
         options = options if options else Options()
 
-        self.service.path = DriverFinder(self.service, options).get_driver_path()
+        self.service.path = self.service.env_path() or DriverFinder(self.service, options).get_driver_path()
         self.service.start()
 
         executor = RemoteConnection(
diff --git a/py/selenium/webdriver/safari/service.py b/py/selenium/webdriver/safari/service.py
@@ -37,15 +37,18 @@ def __init__(
         service_args: typing.Optional[typing.List[str]] = None,
         env: typing.Optional[typing.Mapping[str, str]] = None,
         reuse_service=False,
+        driver_path_env_key: str = None,
         **kwargs,
     ) -> None:
         self.service_args = service_args or []
+        driver_path_env_key = driver_path_env_key or "SE_SAFARIDRIVER"
 
         self.reuse_service = reuse_service
         super().__init__(
             executable_path=executable_path,
             port=port,
             env=env,
+            driver_path_env_key=driver_path_env_key,
             **kwargs,
         )
 
diff --git a/py/selenium/webdriver/safari/webdriver.py b/py/selenium/webdriver/safari/webdriver.py
@@ -45,7 +45,7 @@ def __init__(
         self.service = service if service else Service()
         options = options if options else Options()
 
-        self.service.path = DriverFinder(self.service, options).get_driver_path()
+        self.service.path = self.service.env_path() or DriverFinder(self.service, options).get_driver_path()
 
         if not self.service.reuse_service:
             self.service.start()
diff --git a/py/test/selenium/webdriver/chrome/chrome_service_tests.py b/py/test/selenium/webdriver/chrome/chrome_service_tests.py
@@ -95,3 +95,29 @@ def test_log_output_null_default(driver, capfd) -> None:
     out, err = capfd.readouterr()
     assert "Starting ChromeDriver" not in out
     driver.quit()
+
+
+@pytest.fixture
+def service():
+    return Service()
+
+
+@pytest.mark.usefixtures("service")
+class TestChromeDriverService:
+    service_path = "/path/to/chromedriver"
+
+    @pytest.fixture(autouse=True)
+    def setup_and_teardown(self):
+        os.environ["SE_CHROMEDRIVER"] = self.service_path
+        yield
+        os.environ.pop("SE_CHROMEDRIVER", None)
+
+    def test_uses_path_from_env_variable(self, service):
+        assert "chromedriver" in service.path
+
+    def test_updates_path_after_setting_env_variable(self, service):
+        new_path = "/foo/bar"
+        os.environ["SE_CHROMEDRIVER"] = new_path
+        service.executable_path = self.service_path  # Simulating the update
+
+        assert "chromedriver" in service.executable_path
diff --git a/py/test/selenium/webdriver/firefox/firefox_service_tests.py b/py/test/selenium/webdriver/firefox/firefox_service_tests.py
@@ -17,6 +17,8 @@
 import os
 import subprocess
 
+import pytest
+
 from selenium.webdriver import Firefox
 from selenium.webdriver.firefox.service import Service
 
@@ -54,3 +56,29 @@ def test_log_output_as_stdout(capfd) -> None:
     out, err = capfd.readouterr()
     assert "geckodriver\tINFO\tListening" in out
     driver.quit()
+
+
+@pytest.fixture
+def service():
+    return Service()
+
+
+@pytest.mark.usefixtures("service")
+class TestGeckoDriverService:
+    service_path = "/path/to/geckodriver"
+
+    @pytest.fixture(autouse=True)
+    def setup_and_teardown(self):
+        os.environ["SE_GECKODRIVER"] = self.service_path
+        yield
+        os.environ.pop("SE_GECKODRIVER", None)
+
+    def test_uses_path_from_env_variable(self, service):
+        assert "geckodriver" in service.path
+
+    def test_updates_path_after_setting_env_variable(self, service):
+        new_path = "/foo/bar"
+        os.environ["SE_GECKODRIVER"] = new_path
+        service.executable_path = self.service_path  # Simulating the update
+
+        assert "geckodriver" in service.executable_path
