diff --git a/superset/config.py b/superset/config.py index a6427b528..2d43fc18f 100644 --- a/superset/config.py +++ b/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 diff --git a/superset/utils/webdriver.py b/superset/utils/webdriver.py index e7155ff12..70a4f512e 100644 --- a/superset/utils/webdriver.py +++ b/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 diff --git a/tests/thumbnails_tests.py b/tests/thumbnails_tests.py index 4c0bd4ccb..92d4f9993 100644 --- a/tests/thumbnails_tests.py +++ b/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"