fix: Use page.locator in Playwright reports (#26224)

superset/utils/webdriver.py

@@ -48,8 +48,8 @@
 if feature_flag_manager.is_feature_enabled("PLAYWRIGHT_REPORTS_AND_THUMBNAILS"):
     from playwright.sync_api import (
         BrowserContext,
-        ElementHandle,
         Error as PlaywrightError,
+        Locator,
         Page,
         sync_playwright,
         TimeoutError as PlaywrightTimeout,
@@ -105,14 +105,7 @@ def find_unexpected_errors(page: Page) -> list[str]:
                 alert_div.get_by_role("button").click()
 
                 # wait for modal to show up
-                page.wait_for_selector(
-                    ".ant-modal-content",
-                    timeout=current_app.config[
-                        "SCREENSHOT_WAIT_FOR_ERROR_MODAL_VISIBLE"
-                    ]
-                    * 1000,
-                    state="visible",
-                )
+                page.locator(".ant-modal-content").wait_for(state="visible")
                 err_msg_div = page.locator(".ant-modal-content .ant-modal-body")
                 #
                 # # collect error message
@@ -125,14 +118,7 @@ def find_unexpected_errors(page: Page) -> list[str]:
                 page.locator(".ant-modal-content .ant-modal-close").click()
                 #
                 # # wait until the modal becomes invisible
-                page.wait_for_selector(
-                    ".ant-modal-content",
-                    timeout=current_app.config[
-                        "SCREENSHOT_WAIT_FOR_ERROR_MODAL_INVISIBLE"
-                    ]
-                    * 1000,
-                    state="detached",
-                )
+                page.locator(".ant-modal-content").wait_for(state="detached")
                 try:
                     # Even if some errors can't be updated in the screenshot,
                     # keep all the errors in the server log and do not fail the loop
@@ -147,7 +133,9 @@ def find_unexpected_errors(page: Page) -> list[str]:
 
         return error_messages
 
-    def get_screenshot(self, url: str, element_name: str, user: User) -> bytes | None:
+    def get_screenshot(  # pylint: disable=too-many-locals, too-many-statements
+        self, url: str, element_name: str, user: User
+    ) -> bytes | None:
         with sync_playwright() as playwright:
             browser = playwright.chromium.launch()
             pixel_density = current_app.config["WEBDRIVER_WINDOW"].get(
@@ -166,34 +154,42 @@ def get_screenshot(self, url: str, element_name: str, user: User) -> bytes | Non
             )
             self.auth(user, context)
             page = context.new_page()
-            page.goto(
-                url, wait_until=current_app.config["SCREENSHOT_PLAYWRIGHT_WAIT_EVENT"]
-            )
+            try:
+                page.goto(
+                    url,
+                    wait_until=current_app.config["SCREENSHOT_PLAYWRIGHT_WAIT_EVENT"],
+                )
+            except PlaywrightTimeout:
+                logger.exception(
+                    "Web event %s not detected. Page %s might not have been fully loaded",
+                    current_app.config["SCREENSHOT_PLAYWRIGHT_WAIT_EVENT"],
+                    url,
+                )
+
             img: bytes | None = None
             selenium_headstart = current_app.config["SCREENSHOT_SELENIUM_HEADSTART"]
             logger.debug("Sleeping for %i seconds", selenium_headstart)
             page.wait_for_timeout(selenium_headstart * 1000)
-            element: ElementHandle
+            element: Locator
             try:
                 try:
                     # page didn't load
                     logger.debug(
                         "Wait for the presence of %s at url: %s", element_name, url
                     )
-                    element = page.wait_for_selector(
-                        f".{element_name}",
-                        timeout=self._screenshot_locate_wait * 1000,
-                    )
+                    element = page.locator(f".{element_name}")
+                    element.wait_for()
                 except PlaywrightTimeout as ex:
                     logger.exception("Timed out requesting url %s", url)
                     raise ex
 
                 try:
                     # chart containers didn't render
                     logger.debug("Wait for chart containers to draw at url: %s", url)
-                    page.wait_for_selector(
-                        ".slice_container", timeout=self._screenshot_locate_wait * 1000
-                    )
+                    slice_container_locator = page.locator(".slice_container")
+                    slice_container_locator.first.wait_for()
+                    for slice_container_elem in slice_container_locator.all():
+                        slice_container_elem.wait_for()
                 except PlaywrightTimeout as ex:
                     logger.exception(
                         "Timed out waiting for chart containers to draw at url %s",
@@ -205,11 +201,8 @@ def get_screenshot(self, url: str, element_name: str, user: User) -> bytes | Non
                     logger.debug(
                         "Wait for loading element of charts to be gone at url: %s", url
                     )
-                    page.wait_for_selector(
-                        ".loading",
-                        timeout=self._screenshot_load_wait * 1000,
-                        state="detached",
-                    )
+                    for loading_element in page.locator(".loading").all():
+                        loading_element.wait_for(state="detached")
                 except PlaywrightTimeout as ex:
                     logger.exception(
                         "Timed out waiting for charts to load at url %s", url

superset/views/datasource/utils.py

@@ -43,7 +43,7 @@ def get_limit_clause(page: Optional[int], per_page: Optional[int]) -> dict[str,
     return {"row_offset": offset, "row_limit": limit}
 
 
-def get_samples(  # pylint: disable=too-many-arguments,too-many-locals
+def get_samples(  # pylint: disable=too-many-arguments
     datasource_type: str,
     datasource_id: int,
     force: bool = False,