浏览代码

allow passing kwargs and options to selenium webdriver (#2894)

* allow passing kwargs to selenium webdriver

* always create driver options, add x11 class to chromium

* add driver_option_args
benedikt-bartscher 1 年之前
父节点
当前提交
87a3ddea7f
共有 1 个文件被更改,包括 28 次插入9 次删除
  1. 28 9
      reflex/testing.py

+ 28 - 9
reflex/testing.py

@@ -26,6 +26,7 @@ from typing import (
     AsyncIterator,
     Callable,
     Coroutine,
+    List,
     Optional,
     Type,
     TypeVar,
@@ -513,12 +514,19 @@ class AppHarness:
             raise TimeoutError("Backend is not listening.")
         return backend.servers[0].sockets[0]
 
-    def frontend(self, driver_clz: Optional[Type["WebDriver"]] = None) -> "WebDriver":
+    def frontend(
+        self,
+        driver_clz: Optional[Type["WebDriver"]] = None,
+        driver_kwargs: dict[str, Any] | None = None,
+        driver_option_args: List[str] | None = None,
+    ) -> "WebDriver":
         """Get a selenium webdriver instance pointed at the app.
 
         Args:
             driver_clz: webdriver.Chrome (default), webdriver.Firefox, webdriver.Safari,
                 webdriver.Edge, etc
+            driver_kwargs: additional keyword arguments to pass to the webdriver constructor
+            driver_option_args: additional arguments for the webdriver options
 
         Returns:
             Instance of the given webdriver navigated to the frontend url of the app.
@@ -541,19 +549,30 @@ class AppHarness:
             requested_driver = os.environ.get("APP_HARNESS_DRIVER", "Chrome")
             driver_clz = getattr(webdriver, requested_driver)
             options = getattr(webdriver, f"{requested_driver}Options")()
-        if driver_clz is webdriver.Chrome and want_headless:
+        if driver_clz is webdriver.Chrome:
             options = webdriver.ChromeOptions()
-            options.add_argument("--headless=new")
-        elif driver_clz is webdriver.Firefox and want_headless:
+            options.add_argument("--class=AppHarness")
+            if want_headless:
+                options.add_argument("--headless=new")
+        elif driver_clz is webdriver.Firefox:
             options = webdriver.FirefoxOptions()
-            options.add_argument("-headless")
-        elif driver_clz is webdriver.Edge and want_headless:
+            if want_headless:
+                options.add_argument("-headless")
+        elif driver_clz is webdriver.Edge:
             options = webdriver.EdgeOptions()
-            options.add_argument("headless")
-        if options and (args := os.environ.get("APP_HARNESS_DRIVER_ARGS")):
+            if want_headless:
+                options.add_argument("headless")
+        if options is None:
+            raise RuntimeError(f"Could not determine options for {driver_clz}")
+        if args := os.environ.get("APP_HARNESS_DRIVER_ARGS"):
             for arg in args.split(","):
                 options.add_argument(arg)
-        driver = driver_clz(options=options)  # type: ignore
+        if driver_option_args is not None:
+            for arg in driver_option_args:
+                options.add_argument(arg)
+        if driver_kwargs is None:
+            driver_kwargs = {}
+        driver = driver_clz(options=options, **driver_kwargs)  # type: ignore
         driver.get(self.frontend_url)
         self._frontends.append(driver)
         return driver