Add support for deb (#541)

* add file manager --------- Co-authored-by: s.pirohov <[email protected]>
SergeyPirogov · Jul 23, 2023 · 068aec7 · 068aec7
1 parent f348d4e
commit 068aec7
Show file tree

Hide file tree

Showing 12 changed files with 170 additions and 124 deletions.
diff --git a/README.md b/README.md
@@ -200,6 +200,13 @@ version = read_version_from_cmd("/usr/bin/firefox-bin --version", PATTERN["firef
 driver_binary = FirefoxDriverManager(version=version).install()
 ```
 
+#### Custom Cache and File manager
+
+```python
+cache_manager = DriverCacheManager(file_manager=FileManager())
+manager = ChromeDriverManager(cache_manager=cache_manager)
+```
+
 ## Configuration
 
 **webdriver_manager** has several configuration variables you can be interested in.
@@ -235,15 +242,6 @@ import os
 os.environ['WDM_LOG'] = str(logging.NOTSET)
 ```
 
-### `WDM_PROGRESS_BAR`
-Turn off the progress bar which is displayed on downloads:
-
-```python
-import os
-
-os.environ['WDM_PROGRESS_BAR'] = str(0)
-```
-
 ### `WDM_LOCAL`
 By default, all driver binaries are saved to user.home/.wdm folder. You can override this setting and save binaries to project.root/.wdm.
 

diff --git a/tests/test_downloader.py b/tests/test_downloader.py
@@ -5,27 +5,28 @@
 from webdriver_manager.core.constants import DEFAULT_PROJECT_ROOT_CACHE_PATH
 from webdriver_manager.core.download_manager import WDMDownloadManager
 from webdriver_manager.core.http import WDMHttpClient
-from webdriver_manager.core.utils import save_file, ChromeType
+from webdriver_manager.core.utils import ChromeType, FileManager
 from webdriver_manager.drivers.chrome import ChromeDriver
 
 download_manager = WDMDownloadManager()
+file_manager = FileManager()
 
 
 def test_can_download_driver_as_zip_file(delete_drivers_dir):
     file = download_manager.download_file("http://chromedriver.storage.googleapis.com/2.26/chromedriver_win32.zip")
     assert file.filename == "chromedriver_win32.zip"
-    archive = save_file(file, DEFAULT_PROJECT_ROOT_CACHE_PATH)
+    archive = file_manager.save_archive_file(file, DEFAULT_PROJECT_ROOT_CACHE_PATH)
     assert archive.file_path == f"{DEFAULT_PROJECT_ROOT_CACHE_PATH}{os.sep}{file.filename}"
-    assert archive.unpack(DEFAULT_PROJECT_ROOT_CACHE_PATH) == ["chromedriver.exe"]
+    assert file_manager.unpack_archive(archive, DEFAULT_PROJECT_ROOT_CACHE_PATH) == ["chromedriver.exe"]
 
 
 def test_can_download_driver_as_tar_gz(delete_drivers_dir):
     file = download_manager.download_file(
         "https://github.com/mozilla/geckodriver/releases/download/v0.26.0/geckodriver-v0.26.0-linux32.tar.gz")
     assert file.filename == 'geckodriver-v0.26.0-linux32.tar.gz'
-    archive = save_file(file, DEFAULT_PROJECT_ROOT_CACHE_PATH)
+    archive = file_manager.save_archive_file(file, DEFAULT_PROJECT_ROOT_CACHE_PATH)
     assert archive.file_path == f"{DEFAULT_PROJECT_ROOT_CACHE_PATH}{os.sep}{file.filename}"
-    assert archive.unpack(DEFAULT_PROJECT_ROOT_CACHE_PATH) == ["geckodriver"]
+    assert file_manager.unpack_archive(archive, DEFAULT_PROJECT_ROOT_CACHE_PATH) == ["geckodriver"]
 
 
 @pytest.mark.parametrize('version', ["2.26"])
@@ -37,5 +38,5 @@ def test_can_download_chrome_driver(delete_drivers_dir, version):
 
     file = download_manager.download_file(driver.get_driver_download_url())
     assert file.filename == "chromedriver_win32.zip"
-    archive = save_file(file, DEFAULT_PROJECT_ROOT_CACHE_PATH)
-    assert "chromedriver.exe" in archive.unpack(DEFAULT_PROJECT_ROOT_CACHE_PATH)
+    archive = file_manager.save_archive_file(file, DEFAULT_PROJECT_ROOT_CACHE_PATH)
+    assert "chromedriver.exe" in file_manager.unpack_archive(archive, DEFAULT_PROJECT_ROOT_CACHE_PATH)
diff --git a/tests_xdist/test_cuncurent_1.py b/tests_xdist/test_cuncurent_1.py
@@ -1,10 +1,11 @@
 from selenium import webdriver
+from selenium.webdriver.chrome.service import Service
 
 from webdriver_manager.chrome import ChromeDriverManager
 
 
 def test_chrome_manager_with_selenium():
     driver_path = ChromeDriverManager().install()
-    driver = webdriver.Chrome(driver_path)
+    driver = webdriver.Chrome(service=Service(driver_path))
     driver.get("http://automation-remarks.com")
     driver.close()
diff --git a/tests_xdist/test_cuncurent_2.py b/tests_xdist/test_cuncurent_2.py
@@ -1,10 +1,11 @@
 from selenium import webdriver
+from selenium.webdriver.chrome.service import Service
 
 from webdriver_manager.chrome import ChromeDriverManager
 
 
 def test_chrome_manager_with_selenium():
     driver_path = ChromeDriverManager().install()
-    driver = webdriver.Chrome(driver_path)
+    driver = webdriver.Chrome(service=Service(driver_path))
     driver.get("http://automation-remarks.com")
     driver.close()
diff --git a/webdriver_manager/chrome.py b/webdriver_manager/chrome.py
@@ -2,6 +2,7 @@
 from typing import Optional
 
 from webdriver_manager.core.download_manager import DownloadManager
+from webdriver_manager.core.driver_cache import DriverCacheManager
 from webdriver_manager.core.manager import DriverManager
 from webdriver_manager.core.utils import ChromeType
 from webdriver_manager.drivers.chrome import ChromeDriver
@@ -19,11 +20,13 @@ def __init__(
             chrome_type: str = ChromeType.GOOGLE,
             cache_valid_range: int = 1,
             download_manager: Optional[DownloadManager] = None,
+            cache_manager: Optional[DriverCacheManager] = None
     ):
         super().__init__(
             path,
             cache_valid_range=cache_valid_range,
-            download_manager=download_manager)
+            download_manager=download_manager,
+            cache_manager=cache_manager)
 
         self.driver = ChromeDriver(
             name=name,

diff --git a/webdriver_manager/core/archive.py b/webdriver_manager/core/archive.py
@@ -25,44 +25,4 @@ def __init__(self, path: str, os_type: typing.Optional[str] = None):
         self.file_path = path
         self.os_type: typing.Optional[str] = os_type
 
-    def unpack(self, directory):
-        if self.file_path.endswith(".zip"):
-            return self.__extract_zip(directory)
-        elif self.file_path.endswith(".tar.gz"):
-            return self.__extract_tar_file(directory)
 
-    def __extract_zip(self, to_directory):
-        zip_class = (LinuxZipFileWithPermissions if self.os_type ==
-                     "linux" else zipfile.ZipFile)
-        archive = zip_class(self.file_path)
-        try:
-            archive.extractall(to_directory)
-        except Exception as e:
-            if e.args[0] not in [26, 13] and e.args[1] not in [
-                "Text file busy",
-                "Permission denied",
-            ]:
-                raise e
-            file_names = []
-            for n in archive.namelist():
-                if "/" not in n:
-                    file_names.append(n)
-                else:
-                    file_path, file_name = n.split("/")
-                    full_file_path = os.path.join(to_directory, file_path)
-                    source = os.path.join(full_file_path, file_name)
-                    destination = os.path.join(to_directory, file_name)
-                    os.rename(source, destination)
-                    file_names.append(file_name)
-            return sorted(file_names, key=lambda x: x.lower())
-        return archive.namelist()
-
-    def __extract_tar_file(self, to_directory):
-        try:
-            tar = tarfile.open(self.file_path, mode="r:gz")
-        except tarfile.ReadError:
-            tar = tarfile.open(self.file_path, mode="r:bz2")
-        members = tar.getmembers()
-        tar.extractall(to_directory)
-        tar.close()
-        return [x.name for x in members]
diff --git a/webdriver_manager/core/driver_cache.py b/webdriver_manager/core/driver_cache.py
@@ -9,11 +9,11 @@
 )
 from webdriver_manager.core.driver import Driver
 from webdriver_manager.core.logger import log
-from webdriver_manager.core.utils import get_date_diff, File, save_file
+from webdriver_manager.core.utils import get_date_diff, File, FileManager
 
 
-class DriverCache(object):
-    def __init__(self, root_dir=None, valid_range=1):
+class DriverCacheManager(object):
+    def __init__(self, root_dir=None, valid_range=1, file_manager=None):
         self._root_dir = DEFAULT_USER_HOME_CACHE_PATH
         is_wdm_local = wdm_local()
         xdist_worker_id = get_xdist_worker_id()
@@ -30,15 +30,24 @@ def __init__(self, root_dir=None, valid_range=1):
         self._drivers_json_path = os.path.join(self._root_dir, "drivers.json")
         self._date_format = "%d/%m/%Y"
         self._drivers_directory = os.path.join(self._root_dir, self._drivers_root)
-        self.valid_range = valid_range
+        self._cache_valid_days_range = valid_range
         self._cache_key_driver_version = None
         self._metadata_key = None
         self._driver_binary_path = None
+        self._file_manager = file_manager
+        if not self._file_manager:
+            self._file_manager = FileManager()
+
+    def save_archive_file(self, file: File, path):
+        return self._file_manager.save_archive_file(file, path)
+
+    def unpack_archive(self, archive, path):
+        return self._file_manager.unpack_archive(archive, path)
 
     def save_file_to_cache(self, driver: Driver, file: File):
         path = self.__get_path(driver)
-        archive = save_file(file, path)
-        files = archive.unpack(path)
+        archive = self.save_archive_file(file, path)
+        files = self.unpack_archive(archive, path)
         binary = self.__get_binary(files, driver.get_name())
         binary_path = os.path.join(path, binary)
         self.__save_metadata(driver, binary_path)
@@ -107,7 +116,7 @@ def __is_valid(self, driver_info):
         dates_diff = get_date_diff(
             driver_info["timestamp"], datetime.date.today(), self._date_format
         )
-        return dates_diff < self.valid_range
+        return dates_diff < self._cache_valid_days_range
 
     def load_metadata_content(self):
         if os.path.exists(self._drivers_json_path):

diff --git a/webdriver_manager/core/manager.py b/webdriver_manager/core/manager.py
@@ -1,5 +1,5 @@
 from webdriver_manager.core.download_manager import WDMDownloadManager
-from webdriver_manager.core.driver_cache import DriverCache
+from webdriver_manager.core.driver_cache import DriverCacheManager
 from webdriver_manager.core.logger import log
 
 
@@ -8,8 +8,13 @@ def __init__(
             self,
             root_dir=None,
             cache_valid_range=1,
-            download_manager=None):
-        self.driver_cache = DriverCache(root_dir, cache_valid_range)
+            download_manager=None,
+            cache_manager=None
+    ):
+        self._cache_manager = cache_manager
+        if not self._cache_manager:
+            self._cache_manager = DriverCacheManager(root_dir, cache_valid_range)
+
         self._download_manager = download_manager
         if self._download_manager is None:
             self._download_manager = WDMDownloadManager()
@@ -23,10 +28,10 @@ def install(self) -> str:
         raise NotImplementedError("Please Implement this method")
 
     def _get_driver_path(self, driver):
-        binary_path = self.driver_cache.find_driver(driver)
+        binary_path = self._cache_manager.find_driver(driver)
         if binary_path:
             return binary_path
 
         file = self._download_manager.download_file(driver.get_driver_download_url())
-        binary_path = self.driver_cache.save_file_to_cache(driver, file)
+        binary_path = self._cache_manager.save_file_to_cache(driver, file)
         return binary_path
diff --git a/webdriver_manager/core/utils.py b/webdriver_manager/core/utils.py
@@ -4,8 +4,10 @@
 import re
 import subprocess
 import sys
+import tarfile
+import zipfile
 
-from webdriver_manager.core.archive import Archive
+from webdriver_manager.core.archive import Archive, LinuxZipFileWithPermissions
 
 
 class File(object):
@@ -37,15 +39,59 @@ def filename(self) -> str:
         return filename
 
 
-def save_file(file: File, directory: str):
-    os.makedirs(directory, exist_ok=True)
+class FileManager(object):
 
-    archive_path = f"{directory}{os.sep}{file.filename}"
-    with open(archive_path, "wb") as code:
-        code.write(file.content)
-    if not os.path.exists(archive_path):
-        raise FileExistsError(f"No file has been saved on such path {archive_path}")
-    return Archive(archive_path, os_type=os_name())
+    def save_archive_file(self, file: File, directory: str):
+        os.makedirs(directory, exist_ok=True)
+
+        archive_path = f"{directory}{os.sep}{file.filename}"
+        with open(archive_path, "wb") as code:
+            code.write(file.content)
+        if not os.path.exists(archive_path):
+            raise FileExistsError(f"No file has been saved on such path {archive_path}")
+        return Archive(archive_path, os_type=os_name())
+
+    def unpack_archive(self, archive_file: Archive, target_dir):
+        file_path = archive_file.file_path
+        if file_path.endswith(".zip"):
+            return self.__extract_zip(archive_file, target_dir)
+        elif file_path.endswith(".tar.gz"):
+            return self.__extract_tar_file(archive_file, target_dir)
+
+    def __extract_zip(self, archive_file, to_directory):
+        zip_class = (LinuxZipFileWithPermissions if archive_file.os_type == "linux" else zipfile.ZipFile)
+        archive = zip_class(archive_file.file_path)
+        try:
+            archive.extractall(to_directory)
+        except Exception as e:
+            if e.args[0] not in [26, 13] and e.args[1] not in [
+                "Text file busy",
+                "Permission denied",
+            ]:
+                raise e
+            file_names = []
+            for n in archive.namelist():
+                if "/" not in n:
+                    file_names.append(n)
+                else:
+                    file_path, file_name = n.split("/")
+                    full_file_path = os.path.join(to_directory, file_path)
+                    source = os.path.join(full_file_path, file_name)
+                    destination = os.path.join(to_directory, file_name)
+                    os.rename(source, destination)
+                    file_names.append(file_name)
+            return sorted(file_names, key=lambda x: x.lower())
+        return archive.namelist()
+
+    def __extract_tar_file(self, archive_file, to_directory):
+        try:
+            tar = tarfile.open(archive_file.file_path, mode="r:gz")
+        except tarfile.ReadError:
+            tar = tarfile.open(archive_file.file_path, mode="r:bz2")
+        members = tar.getmembers()
+        tar.extractall(to_directory)
+        tar.close()
+        return [x.name for x in members]
 
 
 class OSType(object):

diff --git a/webdriver_manager/firefox.py b/webdriver_manager/firefox.py
@@ -2,25 +2,30 @@
 from typing import Optional
 
 from webdriver_manager.core.download_manager import DownloadManager
+from webdriver_manager.core.driver_cache import DriverCacheManager
 from webdriver_manager.core.manager import DriverManager
 from webdriver_manager.drivers.firefox import GeckoDriver
 
 
 class GeckoDriverManager(DriverManager):
     def __init__(
-        self,
-        version: Optional[str] = None,
-        os_type: Optional[str] = None,
-        path: Optional[str] = None,
-        name: str = "geckodriver",
-        url: str = "https://github.com/mozilla/geckodriver/releases/download",
-        latest_release_url: str = "https://api.github.com/repos/mozilla/geckodriver/releases/latest",
-        mozila_release_tag: str = "https://api.github.com/repos/mozilla/geckodriver/releases/tags/{0}",
-        cache_valid_range: int = 1,
-        download_manager: Optional[DownloadManager] = None,
+            self,
+            version: Optional[str] = None,
+            os_type: Optional[str] = None,
+            path: Optional[str] = None,
+            name: str = "geckodriver",
+            url: str = "https://github.com/mozilla/geckodriver/releases/download",
+            latest_release_url: str = "https://api.github.com/repos/mozilla/geckodriver/releases/latest",
+            mozila_release_tag: str = "https://api.github.com/repos/mozilla/geckodriver/releases/tags/{0}",
+            cache_valid_range: int = 1,
+            download_manager: Optional[DownloadManager] = None,
+            cache_manager: Optional[DriverCacheManager] = None
     ):
         super(GeckoDriverManager, self).__init__(
-            path, cache_valid_range, download_manager=download_manager
+            path,
+            cache_valid_range,
+            download_manager=download_manager,
+            cache_manager=cache_manager
         )
 
         self.driver = GeckoDriver(