From dba44b1ac1c810f366f3c47faee6206c9c77b31e Mon Sep 17 00:00:00 2001
From: Patrick Robertson <robertson.patrick@gmail.com>
Date: Fri, 7 Mar 2025 12:07:54 +0000
Subject: [PATCH] Use WebDriverWait when waiting for elements in screenshot
 enricher

---
 src/auto_archiver/utils/webdriver.py | 26 +++++++++++++++-----------
 1 file changed, 15 insertions(+), 11 deletions(-)

diff --git a/src/auto_archiver/utils/webdriver.py b/src/auto_archiver/utils/webdriver.py
index c6ad341..1d01df2 100644
--- a/src/auto_archiver/utils/webdriver.py
+++ b/src/auto_archiver/utils/webdriver.py
@@ -9,8 +9,9 @@ from urllib.parse import urlparse, urlunparse
 from http.cookiejar import MozillaCookieJar
 
 from selenium import webdriver
-from selenium.common.exceptions import TimeoutException
-from selenium.webdriver.common.proxy import Proxy, ProxyType
+from selenium.webdriver.support.ui import WebDriverWait
+from selenium.webdriver.support import expected_conditions as EC
+from selenium.common import exceptions as selenium_exceptions
 from selenium.webdriver.common.print_page_options import PrintOptions
 from selenium.webdriver.common.by import By
 
@@ -78,19 +79,22 @@ class CookieSettingDriver(webdriver.Firefox):
         super(CookieSettingDriver, self).get(url)
         if self.facebook_accept_cookies:
             # try and click the 'close' button on the 'login' window to close it
-            close_button = self.find_element(By.XPATH, "//div[@role='dialog']//div[@aria-label='Close']")
-            if close_button:
-                close_button.click()
+            try:
+                xpath = "//div[@role='dialog']//div[@aria-label='Close']"
+                WebDriverWait(self, 5).until(EC.element_to_be_clickable((By.XPATH, xpath))).click()
+            except selenium_exceptions.NoSuchElementException:
+                logger.warning("Unable to find the 'close' button on the facebook login window")
+                pass
+
         else:
 
             # for all other sites, try and use some common button text to reject/accept cookies
             for text in ["Refuse non-essential cookies", "Decline optional cookies", "Reject additional cookies", "Accept all cookies"]:
                 try:
-                    accept_button = self.find_element(By.XPATH, f"//*[contains(translate(text(), 'ABCDEFGHIJKLMNOPQRSTUVWXYZ', 'abcdefghijklmnopqrstuvwxyz'), '{text.lower()}')]")
-                    if accept_button:
-                        accept_button.click()
-                        break
-                except Exception as e:
+                    xpath = f"//*[contains(translate(text(), 'ABCDEFGHIJKLMNOPQRSTUVWXYZ', 'abcdefghijklmnopqrstuvwxyz'), '{text.lower()}')]"
+                    WebDriverWait(self, 5).until(EC.element_to_be_clickable((By.XPATH, xpath))).click()
+                    break
+                except selenium_exceptions.WebDriverException:
                     pass
 
     
@@ -124,7 +128,7 @@ class Webdriver:
             self.driver.set_window_size(self.width, self.height)
             self.driver.set_page_load_timeout(self.timeout_seconds)
             self.driver.print_options = self.print_options
-        except TimeoutException as e:
+        except selenium_exceptions.TimeoutException as e:
             logger.error(f"failed to get new webdriver, possibly due to insufficient system resources or timeout settings: {e}")
 
         return self.driver