diff --git a/src/wabot/hosts.py b/src/wabot/hosts.py new file mode 100644 index 0000000..8b37da7 --- /dev/null +++ b/src/wabot/hosts.py @@ -0,0 +1,88 @@ +"""Driver hosts: the processes that keep a browser alive, and how to reach them. + +Persistence-capable browsers are always driven over the WebDriver wire +protocol (a URL). Two hosts provide that URL: ``ExternalServer`` (a +Selenium Grid or standalone driver the user runs) and ``ManagedService`` +(a chromedriver/geckodriver wabot spawns detached so it outlives Python). +""" + +from __future__ import annotations + +import json +import logging +import urllib.request + +from selenium import webdriver + +LOGGER = logging.getLogger("wabot") + +BROWSER_ALIASES = {"chrome": "chromium"} +DRIVER_BINARIES = {"chromium": "chromedriver", "firefox": "geckodriver"} + + +def normalize_browser(browser: str) -> str: + browser = BROWSER_ALIASES.get(browser, browser) + if browser not in DRIVER_BINARIES: + raise ValueError( + f"unsupported browser {browser!r}; use 'firefox' or 'chromium' (alias: 'chrome')" + ) + return browser + + +def build_options( + browser: str, *, headless: bool = False, user_agent: str | None = None +) -> webdriver.ChromeOptions | webdriver.FirefoxOptions: + """Browser-appropriate Options. UA defaults to the browser's own.""" + browser = normalize_browser(browser) + if browser == "chromium": + opts = webdriver.ChromeOptions() + binary = _chromium_binary() + if binary: + opts.binary_location = binary + if headless: + opts.add_argument("--headless=new") + if user_agent: + opts.add_argument(f"--user-agent={user_agent}") + return opts + opts = webdriver.FirefoxOptions() + if headless: + opts.add_argument("-headless") + if user_agent: + opts.set_preference("general.useragent.override", user_agent) + return opts + + +def _chromium_binary() -> str | None: + import shutil + + for name in ("chromium", "chromium-browser", "google-chrome", "chrome"): + found = shutil.which(name) + if found: + return found + return None + + +def service_alive(url: str, timeout: float = 2.0) -> bool: + """True if a WebDriver server answers GET /status at ``url``. + + Liveness only: geckodriver reports ready=false while its single session + is in use, so the ``ready`` flag is deliberately ignored. + """ + try: + with urllib.request.urlopen(f"{url.rstrip('/')}/status", timeout=timeout) as resp: + json.loads(resp.read()) + return resp.status == 200 + except (OSError, ValueError): + return False + + +class ExternalServer: + """A WebDriver server somebody else runs (Selenium Grid, bare driver).""" + + def __init__(self, url: str): + self.url = url.rstrip("/") + + def ensure_running(self) -> str: + if not service_alive(self.url): + raise ConnectionError(f"no WebDriver server answering at {self.url}/status") + return self.url diff --git a/tests/unit/test_hosts.py b/tests/unit/test_hosts.py new file mode 100644 index 0000000..68025a5 --- /dev/null +++ b/tests/unit/test_hosts.py @@ -0,0 +1,93 @@ +import http.server +import json +import threading + +import pytest + +from wabot.hosts import ExternalServer, build_options, normalize_browser, service_alive + + +class TestNormalizeBrowser: + def test_chrome_is_an_alias_for_chromium(self): + assert normalize_browser("chrome") == "chromium" + + def test_firefox_and_chromium_pass_through(self): + assert normalize_browser("firefox") == "firefox" + assert normalize_browser("chromium") == "chromium" + + def test_unknown_browser_raises(self): + with pytest.raises(ValueError, match="unsupported browser"): + normalize_browser("safari") + + +class TestBuildOptions: + def test_chromium_headless_and_user_agent(self): + opts = build_options("chromium", headless=True, user_agent="MyBot/1.0") + assert "--headless=new" in opts.arguments + assert "--user-agent=MyBot/1.0" in opts.arguments + + def test_chromium_defaults_have_no_headless_or_ua(self): + opts = build_options("chromium") + assert not any(a.startswith(("--headless", "--user-agent")) for a in opts.arguments) + + def test_firefox_headless_and_user_agent(self): + opts = build_options("firefox", headless=True, user_agent="MyBot/1.0") + assert "-headless" in opts.arguments + assert opts.preferences.get("general.useragent.override") == "MyBot/1.0" + + def test_options_type_matches_browser(self): + assert type(build_options("firefox")).__module__.startswith("selenium.webdriver.firefox") + assert type(build_options("chrome")).__module__.startswith("selenium.webdriver.chrome") + + +class _StatusHandler(http.server.BaseHTTPRequestHandler): + ready = True + + def do_GET(self): + if self.path == "/status": + body = json.dumps({"value": {"ready": self.ready, "message": ""}}).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.end_headers() + self.wfile.write(body) + else: + self.send_response(404) + self.end_headers() + + def log_message(self, *args): # keep test output clean + pass + + +@pytest.fixture +def status_server(): + server = http.server.HTTPServer(("127.0.0.1", 0), _StatusHandler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + yield f"http://127.0.0.1:{server.server_address[1]}" + server.shutdown() + + +class TestServiceAlive: + def test_true_when_status_answers(self, status_server): + assert service_alive(status_server) is True + + def test_true_even_when_not_ready(self, status_server): + # geckodriver reports ready=false while its single session is in use; + # aliveness must not require ready==true + _StatusHandler.ready = False + try: + assert service_alive(status_server) is True + finally: + _StatusHandler.ready = True + + def test_false_when_nothing_listens(self): + assert service_alive("http://127.0.0.1:1") is False + + +class TestExternalServer: + def test_returns_url_when_alive(self, status_server): + assert ExternalServer(status_server + "/").ensure_running() == status_server + + def test_raises_when_dead(self): + with pytest.raises(ConnectionError, match="no WebDriver server"): + ExternalServer("http://127.0.0.1:1").ensure_running()