feat: add more timeout configuration on screenshots (#14868)

* feat: more timeout configuration on screenshots

* add tests

superset/config.py

@@ -475,11 +475,17 @@ THUMBNAIL_CACHE_CONFIG: CacheConfig = {
     "CACHE_NO_NULL_WARNING": True,
 }
 
-# Used for thumbnails and other api: Time in seconds before selenium
+# Time in seconds before selenium
 # times out after trying to locate an element on the page and wait
-# for that element to load for an alert screenshot.
+# for that element to load for a screenshot.
 SCREENSHOT_LOCATE_WAIT = 10
+# Time in seconds before selenium
+# times out after waiting for all DOM class elements named "loading" are gone.
 SCREENSHOT_LOAD_WAIT = 60
+# Selenium destroy retries
+SCREENSHOT_SELENIUM_RETRIES = 5
+# Give selenium an headstart, in seconds
+SCREENSHOT_SELENIUM_HEADSTART = 3
 
 # ---------------------------------------------------
 # Image and file configuration

superset/utils/webdriver.py

@@ -16,12 +16,16 @@
 # under the License.
 
 import logging
-import time
+from time import sleep
 from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING
 
 from flask import current_app
 from retry.api import retry_call
-from selenium.common.exceptions import TimeoutException, WebDriverException
+from selenium.common.exceptions import (
+    StaleElementReferenceException,
+    TimeoutException,
+    WebDriverException,
+)
 from selenium.webdriver import chrome, firefox
 from selenium.webdriver.common.by import By
 from selenium.webdriver.remote.webdriver import WebDriver
@@ -33,11 +37,6 @@ from superset.extensions import machine_auth_provider_factory
 WindowSize = Tuple[int, int]
 logger = logging.getLogger(__name__)
 
-# Time in seconds, we will wait for the page to load and render
-SELENIUM_CHECK_INTERVAL = 2
-SELENIUM_RETRIES = 5
-SELENIUM_HEADSTART = 3
-
 
 if TYPE_CHECKING:
     from flask_appbuilder.security.sqla.models import User
@@ -95,18 +94,17 @@ class WebDriverProxy:
             pass
 
     def get_screenshot(
-        self,
-        url: str,
-        element_name: str,
-        user: "User",
-        retries: int = SELENIUM_RETRIES,
+        self, url: str, element_name: str, user: "User",
     ) -> Optional[bytes]:
+
         driver = self.auth(user)
         driver.set_window_size(*self._window)
         driver.get(url)
         img: Optional[bytes] = None
-        logger.debug("Sleeping for %i seconds", SELENIUM_HEADSTART)
-        time.sleep(SELENIUM_HEADSTART)
+        selenium_headstart = current_app.config["SCREENSHOT_SELENIUM_HEADSTART"]
+        logger.debug("Sleeping for %i seconds", selenium_headstart)
+        sleep(selenium_headstart)
+
         try:
             logger.debug("Wait for the presence of %s", element_name)
             element = WebDriverWait(driver, self._screenshot_locate_wait).until(
@@ -120,11 +118,14 @@ class WebDriverProxy:
             img = element.screenshot_as_png
         except TimeoutException:
             logger.error("Selenium timed out requesting url %s", url, exc_info=True)
+        except StaleElementReferenceException:
+            logger.error(
+                "Selenium timed out while waiting for chart(s) to load %s",
+                url,
+                exc_info=True,
+            )
         except WebDriverException as ex:
             logger.error(ex, exc_info=True)
-            # Some webdrivers do not support screenshots for elements.
-            # In such cases, take a screenshot of the entire page.
-            img = driver.screenshot()  # pylint: disable=no-member
         finally:
-            self.destroy(driver, retries)
+            self.destroy(driver, current_app.config["SCREENSHOT_SELENIUM_RETRIES"])
         return img

tests/thumbnails_tests.py

@@ -19,7 +19,7 @@
 import urllib.request
 from io import BytesIO
 from unittest import skipUnless
-from unittest.mock import patch
+from unittest.mock import ANY, call, patch
 
 from flask_testing import LiveServerTestCase
 from sqlalchemy.sql import func
@@ -29,7 +29,8 @@ from superset.extensions import machine_auth_provider_factory
 from superset.models.dashboard import Dashboard
 from superset.models.slice import Slice
 from superset.utils.screenshots import ChartScreenshot, DashboardScreenshot
-from superset.utils.urls import get_url_host
+from superset.utils.urls import get_url_host, get_url_path
+from superset.utils.webdriver import WebDriverProxy
 from tests.conftest import with_feature_flags
 from tests.test_app import app
 
@@ -61,6 +62,47 @@ class TestThumbnailsSeleniumLive(LiveServerTestCase):
             self.assertEqual(response.getcode(), 202)
 
 
+class TestWebDriverProxy(SupersetTestCase):
+    @patch("superset.utils.webdriver.WebDriverWait")
+    @patch("superset.utils.webdriver.firefox")
+    @patch("superset.utils.webdriver.sleep")
+    def test_screenshot_selenium_headstart(
+        self, mock_sleep, mock_webdriver, mock_webdriver_wait
+    ):
+        webdriver = WebDriverProxy("firefox")
+        user = security_manager.get_user_by_username(
+            app.config["THUMBNAIL_SELENIUM_USER"]
+        )
+        url = get_url_path("Superset.slice", slice_id=1, standalone="true")
+        app.config["SCREENSHOT_SELENIUM_HEADSTART"] = 5
+        webdriver.get_screenshot(url, "chart-container", user=user)
+        assert mock_sleep.call_args_list[0] == call(5)
+
+    @patch("superset.utils.webdriver.WebDriverWait")
+    @patch("superset.utils.webdriver.firefox")
+    def test_screenshot_selenium_locate_wait(self, mock_webdriver, mock_webdriver_wait):
+        app.config["SCREENSHOT_LOCATE_WAIT"] = 15
+        webdriver = WebDriverProxy("firefox")
+        user = security_manager.get_user_by_username(
+            app.config["THUMBNAIL_SELENIUM_USER"]
+        )
+        url = get_url_path("Superset.slice", slice_id=1, standalone="true")
+        webdriver.get_screenshot(url, "chart-container", user=user)
+        assert mock_webdriver_wait.call_args_list[0] == call(ANY, 15)
+
+    @patch("superset.utils.webdriver.WebDriverWait")
+    @patch("superset.utils.webdriver.firefox")
+    def test_screenshot_selenium_load_wait(self, mock_webdriver, mock_webdriver_wait):
+        app.config["SCREENSHOT_LOAD_WAIT"] = 15
+        webdriver = WebDriverProxy("firefox")
+        user = security_manager.get_user_by_username(
+            app.config["THUMBNAIL_SELENIUM_USER"]
+        )
+        url = get_url_path("Superset.slice", slice_id=1, standalone="true")
+        webdriver.get_screenshot(url, "chart-container", user=user)
+        assert mock_webdriver_wait.call_args_list[1] == call(ANY, 15)
+
+
 class TestThumbnails(SupersetTestCase):
 
     mock_image = b"bytes mock image"