Skip to content

Commit

Permalink
fix(python/webdriver): make console.log handling more robust
Browse files Browse the repository at this point in the history 
Replace the use of the low-level Chrome DevTools protocol APIs with the
newer bidi APIs.
  • Loading branch information
@jbms
jbms committed Feb 21, 2025
1 parent e72d996 commit 9f18e2d
Show file tree
Hide file tree
Showing 2 changed files with 51 additions and 92 deletions.
100 changes: 27 additions & 73 deletions python/neuroglancer/webdriver.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,29 +18,28 @@
import sys
import threading
import time
from collections.abc import Sequence
from typing import Callable, NamedTuple, Optional
from collections.abc import Callable, Sequence
from typing import TYPE_CHECKING

if TYPE_CHECKING:
from selenium.webdriver.common.bidi.script import ConsoleLogEntry
else:
ConsoleLogEntry = None

class LogMessage(NamedTuple):
message: str
level: Optional[str]


LogListener = Callable[[LogMessage], None]
LogListener = Callable[[ConsoleLogEntry], None]


class WebdriverBase:
def __init__(
self,
headless=True,
browser="chrome",
browser_binary_path: Optional[str] = None,
browser_binary_path: str | None = None,
window_size=(1920, 1080),
debug=False,
docker=False,
print_logs=True,
extra_command_line_args: Optional[Sequence[str]] = None,
extra_command_line_args: Sequence[str] | None = None,
):
self.headless = headless
self.browser = browser
Expand All @@ -52,23 +51,20 @@ def __init__(
)
self.debug = debug
self.browser_binary_path = browser_binary_path
self._log_listeners_lock = threading.Lock()
self._log_listeners: dict[LogListener, None] = {}

self._closed = False
self._init_driver()

if print_logs:
self.add_log_listener(
lambda log: print(
f"console.{log.level}: {log.message}", file=sys.stderr
)
lambda log: print(f"console.{log.level}: {log.text}", file=sys.stderr)
)

self._closed = False
self._init_driver()

def _init_chrome(self):
import selenium.webdriver

chrome_options = selenium.webdriver.ChromeOptions()
chrome_options.enable_bidi = True
if self.headless:
chrome_options.add_argument("--headless=new")
chrome_options.add_experimental_option("excludeSwitches", ["enable-automation"])
Expand Down Expand Up @@ -96,13 +92,12 @@ def _init_firefox(self):
options.arguments.extend(self.extra_command_line_args)
if self.browser_binary_path:
options.binary_location = self.browser_binary_path
options.enable_bidi = True
self.driver = selenium.webdriver.Firefox(
options=options,
)

def _init_driver(self):
import trio

if self.browser == "chrome":
self._init_chrome()
elif self.browser == "firefox":
Expand All @@ -112,80 +107,39 @@ def _init_driver(self):
f'unsupported browser: {self.browser}, must be "chrome" or "firefox"'
)

def log_handler(driver):
async def start_listening(listener):
async for event in listener:
message = LogMessage(message=event.args[0].value, level=event.type_)
with self._log_listeners_lock:
for listener in self._log_listeners:
listener(message)

async def start_listening_for_exceptions(listener):
async for event in listener:
message = LogMessage(
message=event.exception_details.text, level="exception"
)
with self._log_listeners_lock:
for listener in self._log_listeners:
listener(message)

async def run():
async with self.driver.bidi_connection() as connection:
session, devtools = connection.session, connection.devtools
await session.execute(devtools.page.enable())
await session.execute(devtools.runtime.enable())
listener = session.listen(devtools.runtime.ConsoleAPICalled)
exception_listener = session.listen(
devtools.runtime.ExceptionThrown
)
with trio.CancelScope() as cancel_scope:
async with trio.open_nursery() as nursery:
nursery.start_soon(start_listening, listener)
nursery.start_soon(
start_listening_for_exceptions, exception_listener
)
while True:
await trio.sleep(2)
if not driver.service.is_connectable():
cancel_scope.cancel()

trio.run(run)

t = threading.Thread(target=log_handler, args=(self.driver,))
t.daemon = True
t.start()

def __enter__(self):
return self

def __exit__(self, exc_type, exc_value, traceback):
self.driver.quit()
self._closed = True

def add_log_listener(self, listener: LogListener):
with self._log_listeners_lock:
self._log_listeners[listener] = None
def add_log_listener(self, listener: LogListener) -> Callable[[], None]:
console_id = self.driver.script.add_console_message_handler(listener)
error_id = self.driver.script.add_javascript_error_handler(listener)

def unregister():
self.driver.script.remove_console_message_handler(console_id)
self.driver.script.remove_javascript_error_handler(error_id)

def remove_log_listener(self, listener: LogListener):
with self._log_listeners_lock:
return self._log_listeners.pop(listener, True) is None
return unregister

@contextlib.contextmanager
def log_listener(self, listener: LogListener):
unregister = self.add_log_listener(listener)
try:
self.add_log_listener(listener)
yield
finally:
self.remove_log_listener(listener)
unregister()

@contextlib.contextmanager
def wait_for_log_message(self, pattern: str, timeout: Optional[float] = None):
def wait_for_log_message(self, pattern: str, timeout: float | None = None):
event = threading.Event()

def handle_message(msg):
if event.is_set():
return
if re.fullmatch(pattern, msg.message):
if re.fullmatch(pattern, msg.text):
event.set()

with self.log_listener(handle_message):
Expand Down
43 changes: 24 additions & 19 deletions python/tests/conftest.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,11 +11,11 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import atexit

import os
import pathlib
from collections.abc import Iterator
from typing import Callable
import threading
from collections.abc import Callable, Iterator

import neuroglancer.static_file_server
import neuroglancer.webdriver
Expand Down Expand Up @@ -78,6 +78,25 @@ def pytest_addoption(parser):
)


def _setup_webdriver(request, cls):
webdriver = cls(
headless=request.config.getoption("--headless"),
docker=request.config.getoption("--webdriver-docker"),
debug=request.config.getoption("--debug-webdriver"),
browser=request.config.getoption("--browser"),
browser_binary_path=request.config.getoption("--browser-binary-path"),
)

# Note: Regular atexit functions are run only after non-daemon threads are joined.
# However, Selenium creates a non-daemon thread for bidi websocket communication,
# which blocks Python from exiting.
#
# The `threading._register_atexit` function registers an early atexit callback to be
# invoked *before* non-daemon threads are joined.
threading._register_atexit(webdriver.driver.quit)
return webdriver


@pytest.fixture(scope="session")
def _webdriver_internal(request):
if request.config.getoption("--skip-browser-tests"):
Expand All @@ -88,31 +107,17 @@ def _webdriver_internal(request):
static_content_url = request.config.getoption("--static-content-url")
if static_content_url is not None:
neuroglancer.set_static_content_source(url=static_content_url)
webdriver = neuroglancer.webdriver.Webdriver(
headless=request.config.getoption("--headless"),
docker=request.config.getoption("--webdriver-docker"),
debug=request.config.getoption("--debug-webdriver"),
browser=request.config.getoption("--browser"),
browser_binary_path=request.config.getoption("--browser-binary-path"),
)
webdriver = _setup_webdriver(request, neuroglancer.webdriver.Webdriver)
if request.config.getoption("--neuroglancer-server-debug"):
neuroglancer.server.debug = True
atexit.register(webdriver.driver.close)
return webdriver


@pytest.fixture(scope="session")
def webdriver_generic(request):
if request.config.getoption("--skip-browser-tests"):
pytest.skip("--skip-browser-tests")
webdriver = neuroglancer.webdriver.WebdriverBase(
headless=request.config.getoption("--headless"),
docker=request.config.getoption("--webdriver-docker"),
debug=request.config.getoption("--debug-webdriver"),
browser=request.config.getoption("--browser"),
browser_binary_path=request.config.getoption("--browser-binary-path"),
)
atexit.register(webdriver.driver.close)
webdriver = _setup_webdriver(request, neuroglancer.webdriver.WebdriverBase)
return webdriver


Expand Down

0 comments on commit 9f18e2d

Please sign in to comment.