From 989fe27a0f266d15af63cee41ab95cfcba9db493 Mon Sep 17 00:00:00 2001
From: Elizabeth Thompson <eschutho@gmail.com>
Date: Thu, 9 Mar 2023 12:55:12 -0800
Subject: [PATCH] chore: log different selenium timeout errors differently
 (#23290)

---
 superset/utils/webdriver.py | 54 ++++++++++++++++++++++++++-----------
 1 file changed, 39 insertions(+), 15 deletions(-)

diff --git a/superset/utils/webdriver.py b/superset/utils/webdriver.py
index 92a73d347c20a..05dbee674052d 100644
--- a/superset/utils/webdriver.py
+++ b/superset/utils/webdriver.py
@@ -176,22 +176,46 @@ def get_screenshot(
         sleep(selenium_headstart)
 
         try:
-            logger.debug("Wait for the presence of %s", element_name)
-            element = WebDriverWait(driver, self._screenshot_locate_wait).until(
-                EC.presence_of_element_located((By.CLASS_NAME, element_name))
-            )
+            try:
+                # page didn't load
+                logger.debug(
+                    "Wait for the presence of %s at url: %s", element_name, url
+                )
+                element = WebDriverWait(driver, self._screenshot_locate_wait).until(
+                    EC.presence_of_element_located((By.CLASS_NAME, element_name))
+                )
+            except TimeoutException as ex:
+                logger.exception("Selenium timed out requesting url %s", url)
+                raise ex
 
-            logger.debug("Wait for chart containers to draw")
-            WebDriverWait(driver, self._screenshot_locate_wait).until(
-                EC.visibility_of_all_elements_located(
-                    (By.CLASS_NAME, "slice_container")
+            try:
+                # chart containers didn't render
+                logger.debug("Wait for chart containers to draw at url: %s", url)
+                WebDriverWait(driver, self._screenshot_locate_wait).until(
+                    EC.visibility_of_all_elements_located(
+                        (By.CLASS_NAME, "slice_container")
+                    )
                 )
-            )
+            except TimeoutException as ex:
+                logger.exception(
+                    "Selenium timed out waiting for chart containers to draw at url %s",
+                    url,
+                )
+                raise ex
 
-            logger.debug("Wait for loading element of charts to be gone")
-            WebDriverWait(driver, self._screenshot_load_wait).until_not(
-                EC.presence_of_all_elements_located((By.CLASS_NAME, "loading"))
-            )
+            try:
+                # charts took too long to load
+                logger.debug(
+                    "Wait for loading element of charts to be gone at url: %s", url
+                )
+                WebDriverWait(driver, self._screenshot_load_wait).until_not(
+                    EC.presence_of_all_elements_located((By.CLASS_NAME, "loading"))
+                )
+            except TimeoutException as ex:
+                logger.exception(
+                    "Selenium timed out waiting for charts to load at url %s", url
+                )
+                raise ex
 
             selenium_animation_wait = current_app.config[
                 "SCREENSHOT_SELENIUM_ANIMATION_WAIT"
@@ -215,9 +239,9 @@ def get_screenshot(
                     )
 
             img = element.screenshot_as_png
-
         except TimeoutException:
-            logger.exception("Selenium timed out requesting url %s", url)
+            # raise again for the finally block, but handled above
+            pass
         except StaleElementReferenceException:
             logger.exception(
                 "Selenium got a stale element while requesting url %s",