Merge pull request #10 from hhslepicka/wait_for_downloads

ENH: Add wait_for_downloads method.

Author: hhslepicka

botcity/web/bot.py

@@ -17,6 +17,7 @@
 from selenium.common.exceptions import InvalidSessionIdException
 from selenium.webdriver.common.action_chains import ActionChains
 from selenium.webdriver.common.keys import Keys
+from selenium.webdriver.support.ui import WebDriverWait
 
 from . import config, cv2find
 from .browsers import Browser, BROWSER_CONFIGS
@@ -65,7 +66,7 @@ def __init__(self, headless=False):
         self._shift_hold = False
 
         self._dimensions = (1600, 900)
-        self._download_folder_path = None  # Defaults to ~/Desktop
+        self._download_folder_path = os.path.join(os.path.expanduser("~"), "Desktop")
 
     @property
     def driver(self):
@@ -205,6 +206,7 @@ def stop_browser(self):
         """
         if not self._driver:
             return
+        self._driver.close()
         self._driver.quit()
         self._driver = None
 
@@ -854,6 +856,17 @@ def print_pdf(self, path=None, print_options=None):
             f.write(bytes_file)
         return path
 
+    def wait_for_downloads(self, timeout: int = 120000):
+        """
+        Wait for all downloads to be finished.
+
+        Args:
+            timeout (int, optional): Timeout in millis. Defaults to 120000.
+        """
+        wait_method = BROWSER_CONFIGS.get(self.browser).get("wait_for_downloads")
+        # waits for all the files to be completed
+        WebDriverWait(self._driver, timeout/1000, 1).until(wait_method)
+
     #######
     # Mouse
     #######

botcity/web/browsers/__init__.py

@@ -23,16 +23,19 @@ class Browser(str, enum.Enum):
     Browser.CHROME: {
         "driver": "chromedriver",
         "class": chrome.Chrome,
-        "options": chrome.default_options
+        "options": chrome.default_options,
+        "wait_for_downloads": chrome.wait_for_downloads
     },
     Browser.FIREFOX: {
         "driver": "geckodriver",
         "class": firefox.Firefox,
-        "options": firefox.default_options
+        "options": firefox.default_options,
+        "wait_for_downloads": firefox.wait_for_downloads
     },
     Browser.EDGE: {
         "driver": "msedgedriver",
         "class": edge.Edge,
-        "options": edge.default_options
+        "options": edge.default_options,
+        "wait_for_downloads": edge.wait_for_downloads
     },
 }

botcity/web/browsers/chrome.py

@@ -78,3 +78,14 @@ def default_options(headless=False, download_folder_path=None, user_data_dir=Non
     chrome_options.add_argument("--kiosk-printing")
 
     return chrome_options
+
+
+def wait_for_downloads(driver):
+    if not driver.current_url.startswith("chrome://downloads"):
+        driver.get("chrome://downloads/")
+    return driver.execute_script("""
+        var items = document.querySelector('downloads-manager')
+            .shadowRoot.getElementById('downloadsList').items;
+        if (items.every(e => e.state === "COMPLETE"))
+            return items.map(e => e.fileUrl || e.file_url);
+        """)

botcity/web/browsers/edge.py

@@ -83,3 +83,14 @@ def default_options(headless=False, download_folder_path=None, user_data_dir=Non
     edge_options.add_argument("--kiosk-printing")
 
     return edge_options
+
+
+def wait_for_downloads(driver):
+    if not driver.current_url.startswith("chrome://downloads"):
+        driver.get("chrome://downloads/")
+    return driver.execute_script("""
+        var items = document.querySelector('downloads-manager')
+            .shadowRoot.getElementById('downloadsList').items;
+        if (items.every(e => e.state === "COMPLETE"))
+            return items.map(e => e.fileUrl || e.file_url);
+        """)

botcity/web/browsers/firefox.py

@@ -351,9 +351,24 @@ def default_options(headless=False, download_folder_path=None, user_data_dir=Non
     firefox_profile.set_preference('general.warnOnAboutConfig', False)
 
     mimetypes_to_download = ",".join(FIREFOX_MIMETYPES_TO_DOWNLOAD)
+    firefox_profile.set_preference("pdfjs.disabled", True)
+    firefox_profile.set_preference("plugin.disable_full_page_plugin_for_types", mimetypes_to_download)
     firefox_profile.set_preference('browser.helperApps.neverAsk.saveToDisk', mimetypes_to_download)
 
     firefox_profile.update_preferences()
     firefox_options.profile = firefox_profile
 
     return firefox_options
+
+
+def wait_for_downloads(driver):
+    print('Start Wait for Downloads')
+    if not driver.current_url.startswith("about:downloads"):
+        print('Open the downloads tab')
+        driver.get("about:downloads")
+
+    return driver.execute_script("""
+        var items = document.querySelector('richlistbox').itemChildren;
+        if (items.every(e => e.attributes["state"].value != 0))
+            return true;
+        """)