Refactor browser fixtures for tests

Author: stephanlensky

tests/browser_args.py

@@ -1,14 +1,14 @@
 import logging
 import os
 import signal
+from contextlib import AbstractAsyncContextManager
 from enum import Enum
 from threading import Event
 from typing import AsyncGenerator
 
 import pytest
 
 import zendriver as zd
-from tests.browser_args import get_browser_args
 
 logger = logging.getLogger(__name__)
 
@@ -28,26 +28,82 @@ def fixture_params(self):
             return [{"headless": True}, {"headless": False}]
 
 
-BROWSER_MODE = BrowserMode(os.getenv("ZENDRIVER_TEST_BROWSERS", "all"))
-
-PAUSE_AFTER_TEST = os.getenv("ZENDRIVER_PAUSE_AFTER_TEST", "false") == "true"
 NEXT_TEST_EVENT = Event()
 
 
-@pytest.fixture(params=BROWSER_MODE.fixture_params)
-async def browser(request: pytest.FixtureRequest) -> AsyncGenerator[zd.Browser, None]:
+class TestConfig:
+    BROWSER_MODE = BrowserMode(os.getenv("ZENDRIVER_TEST_BROWSERS", "all"))
+    PAUSE_AFTER_TEST = os.getenv("ZENDRIVER_PAUSE_AFTER_TEST", "false") == "true"
+    SANDBOX = os.getenv("ZENDRIVER_TEST_SANDBOX", "false") == "true"
+    USE_WAYLAND = os.getenv("WAYLAND_DISPLAY") is not None
+
+
+class CreateBrowser(AbstractAsyncContextManager):
+    def __init__(
+        self,
+        *,
+        headless: bool = True,
+        sandbox: bool = TestConfig.SANDBOX,
+        browser_args: list[str] | None = None,
+    ):
+        self.headless = headless
+        self.sandbox = sandbox
+        self.browser_args = browser_args
+        self.browser: zd.Browser | None = None
+        self.browser_pid: int | None = None
+
+    def _browser_args(self) -> list[str]:
+        args = []
+        if not self.headless and TestConfig.USE_WAYLAND:
+            # use wayland backend instead of x11
+            args.extend(
+                ["--disable-features=UseOzonePlatform", "--ozone-platform=wayland"]
+            )
+        if self.browser_args is not None:
+            args.extend(self.browser_args)
+
+        return args
+
+    def config(self) -> zd.Config:
+        return zd.Config(
+            headless=self.headless,
+            sandbox=self.sandbox,
+            browser_args=self._browser_args(),
+        )
+
+    async def __aenter__(self) -> zd.Browser:
+        self.browser = await zd.start(self.config())
+        browser_pid = self.browser._process_pid
+        assert browser_pid is not None and browser_pid > 0
+        await self.browser.wait(0)
+        return self.browser
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        if self.browser is not None:
+            await self.browser.stop()
+            assert self.browser_pid is None
+
+
+@pytest.fixture
+def create_browser() -> type[CreateBrowser]:
+    return CreateBrowser
+
+
+@pytest.fixture(params=TestConfig.BROWSER_MODE.fixture_params)
+def headless(request: pytest.FixtureRequest) -> bool:
+    return request.param["headless"]
+
+
+@pytest.fixture
+async def browser(
+    headless: bool, create_browser: type[CreateBrowser]
+) -> AsyncGenerator[zd.Browser, None]:
     NEXT_TEST_EVENT.clear()
-    args = get_browser_args(headless=request.param["headless"])
-    browser = await zd.start(
-        # use wayland for rendering instead of default X11 backend
-        browser_args=args,
-        headless=request.param["headless"],
-    )
-    browser_pid = browser._process_pid
-    assert browser_pid is not None and browser_pid > 0
-    await browser.wait(0)
-    yield browser
-    if PAUSE_AFTER_TEST:
+
+    async with create_browser(headless=headless) as browser:
+        yield browser
+
+    if TestConfig.PAUSE_AFTER_TEST:
         logger.info(
             "Pausing after test. Send next test hotkey (default Mod+Return) to continue to next test"
         )
@@ -58,7 +114,7 @@ async def browser(request: pytest.FixtureRequest) -> AsyncGenerator[zd.Browser,
 
 # signal handler for starting next test
 def handle_next_test(signum, frame):
-    if not PAUSE_AFTER_TEST:
+    if not TestConfig.PAUSE_AFTER_TEST:
         logger.warning(
             "Next test signal received, but ZENDRIVER_PAUSE_AFTER_TEST is not set."
         )

tests/conftest.py

@@ -2,23 +2,22 @@
 from pytest_mock import MockerFixture
 
 import zendriver as zd
-from tests.browser_args import get_browser_args
+from tests.conftest import CreateBrowser
 
 
 async def test_connection_error_raises_exception_and_logs_stderr(
-    mocker: MockerFixture, caplog: pytest.LogCaptureFixture
+    create_browser: type[CreateBrowser],
+    mocker: MockerFixture,
+    caplog: pytest.LogCaptureFixture,
 ):
     mocker.patch(
         "zendriver.core.browser.Browser.test_connection",
         return_value=False,
     )
     with caplog.at_level("INFO"):
         with pytest.raises(Exception):
-            await zd.start(
-                headless=True,
-                browser_args=get_browser_args(headless=True),
-                browser_connection_max_tries=1,
-            )
+            async with create_browser() as _:
+                pass
     assert "Browser stderr" in caplog.text
 
 

tests/core/test_browser.py

@@ -1,18 +1,9 @@
 import zendriver as zd
-from tests.browser_args import get_browser_args
+from tests.conftest import CreateBrowser
 
-# browser fixture not used here because this is a special case
 
-
-async def test_multiple_browsers_diff_userdata():
-    config = zd.Config(
-        headless=True,
-        browser_args=get_browser_args(headless=True)
-        + [
-            "--window-size=800,800",
-            "--wait-for-debugger-webui",
-        ],
-    )
+async def test_multiple_browsers_diff_userdata(create_browser: type[CreateBrowser]):
+    config = create_browser().config()
     config.browser_connection_timeout = 1
     config.browser_connection_max_tries = 15