diff --git a/src/wabot/hosts.py b/src/wabot/hosts.py index b6b70f9..16f5c90 100644 --- a/src/wabot/hosts.py +++ b/src/wabot/hosts.py @@ -22,6 +22,9 @@ from pathlib import Path import platformdirs from selenium import webdriver +from selenium.common.exceptions import WebDriverException +from selenium.webdriver.remote.webdriver import WebDriver as RemoteWebDriver +from urllib3.exceptions import HTTPError as _Urllib3HTTPError LOGGER = logging.getLogger("wabot") @@ -217,3 +220,38 @@ def stop_service(pid: int) -> bool: break time.sleep(0.05) return True + + +class ReattachingRemote(RemoteWebDriver): + """A Remote driver that adopts an existing session instead of creating one. + + Selenium has no supported reattach API; this relies on two internals + that have been stable across 4.x (pin selenium <5): ``execute()`` + injects ``self.session_id`` into every command, and ``start_session`` + is the only place a NEW_SESSION request happens. + """ + + def __init__(self, command_executor: str, session_id: str, options): + # NOT self.session_id: the base __init__ resets that to None before + # calling start_session, so stash the id under a private name. + self._reattach_session_id = session_id + super().__init__(command_executor=command_executor, options=options) + + def start_session(self, capabilities: dict) -> None: + # Skip Command.NEW_SESSION entirely; adopt the saved session. + self.session_id = self._reattach_session_id + self.caps = capabilities # requested caps: fine for normal commands (no CDP/BiDi) + + +def attach(executor_url: str, session_id: str, browser: str): + """Return a live driver adopting the saved session, or None if it is dead.""" + driver = ReattachingRemote(executor_url, session_id, options=build_options(browser)) + try: + _ = driver.current_url # first real HTTP call; raises if the session is gone + return driver + except WebDriverException as ex: + LOGGER.warning("saved session %s is dead: %s", session_id, ex) + return None + except (OSError, _Urllib3HTTPError) as ex: # server unreachable (connection refused, timeout) + LOGGER.warning("no server at %s: %s", executor_url, ex) + return None diff --git a/tests/unit/test_reattach.py b/tests/unit/test_reattach.py new file mode 100644 index 0000000..ca677a5 --- /dev/null +++ b/tests/unit/test_reattach.py @@ -0,0 +1,41 @@ +from selenium.common.exceptions import WebDriverException + +from wabot.hosts import ReattachingRemote, attach, build_options + + +class TestReattachingRemote: + def test_adopts_session_id_without_creating_a_session(self): + driver = ReattachingRemote( + "http://127.0.0.1:1", "saved-session-id", options=build_options("chromium") + ) + # base __init__ resets session_id to None BEFORE start_session runs; + # our override must still win + assert driver.session_id == "saved-session-id" + assert isinstance(driver.caps, dict) + + def test_works_for_firefox_options_too(self): + driver = ReattachingRemote( + "http://127.0.0.1:1", "sid", options=build_options("firefox") + ) + assert driver.session_id == "sid" + + +class TestAttach: + def test_returns_driver_when_probe_succeeds(self, monkeypatch): + monkeypatch.setattr( + ReattachingRemote, "current_url", property(lambda self: "http://example.com") + ) + driver = attach("http://127.0.0.1:1", "sid", "chromium") + assert driver is not None + assert driver.session_id == "sid" + + def test_returns_none_when_session_is_dead(self, monkeypatch): + def boom(self): + raise WebDriverException("invalid session id") + + monkeypatch.setattr(ReattachingRemote, "current_url", property(boom)) + assert attach("http://127.0.0.1:1", "sid", "chromium") is None + + def test_returns_none_when_server_is_unreachable(self): + # no server on port 1: connection refused is immediate; attach must not raise + assert attach("http://127.0.0.1:1", "sid", "chromium") is None