fix: Add waiting time for chart animation when screenshot (apache#15610)

CybercentreCanada · Dec 17, 2021 · 5765c34 · 5765c34
1 parent 60e4019
commit 5765c34
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 0 deletions.
diff --git a/superset/config.py b/superset/config.py
@@ -499,6 +499,8 @@ def _try_json_readsha(  # pylint: disable=unused-argument
 SCREENSHOT_SELENIUM_RETRIES = 5
 # Give selenium an headstart, in seconds
 SCREENSHOT_SELENIUM_HEADSTART = 3
+# Wait for the chart animation, in seconds
+SCREENSHOT_SELENIUM_ANIMATION_WAIT = 5
 
 # ---------------------------------------------------
 # Image and file configuration

diff --git a/superset/utils/webdriver.py b/superset/utils/webdriver.py
@@ -120,6 +120,11 @@ def get_screenshot(
                     (By.CLASS_NAME, "slice_container")
                 )
             )
+            selenium_animation_wait = current_app.config[
+                "SCREENSHOT_SELENIUM_ANIMATION_WAIT"
+            ]
+            logger.debug("Wait %i seconds for chart animation", selenium_animation_wait)
+            sleep(selenium_animation_wait)
             logger.info("Taking a PNG screenshot or url %s", url)
             img = element.screenshot_as_png
         except TimeoutException:

diff --git a/tests/integration_tests/thumbnails_tests.py b/tests/integration_tests/thumbnails_tests.py
@@ -102,6 +102,21 @@ def test_screenshot_selenium_load_wait(self, mock_webdriver, mock_webdriver_wait
         webdriver.get_screenshot(url, "chart-container", user=user)
         assert mock_webdriver_wait.call_args_list[1] == call(ANY, 15)
 
+    @patch("superset.utils.webdriver.WebDriverWait")
+    @patch("superset.utils.webdriver.firefox")
+    @patch("superset.utils.webdriver.sleep")
+    def test_screenshot_selenium_animation_wait(
+        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_ANIMATION_WAIT"] = 4
+        webdriver.get_screenshot(url, "chart-container", user=user)
+        assert mock_sleep.call_args_list[1] == call(4)
+
 
 class TestThumbnails(SupersetTestCase):