Merge pull request #100 from kayqueGovetri/0.branch

FIX: Implement atexit to stop browser

Author: hhslepicka

botcity/web/bot.py

@@ -1,3 +1,6 @@
+from __future__ import annotations
+
+import atexit
 import base64
 import functools
 import glob
@@ -25,6 +28,7 @@
 from selenium.webdriver.remote.webelement import WebElement
 from selenium.webdriver.support.wait import WebDriverWait, TimeoutException, NoSuchElementException
 from selenium.webdriver.support import expected_conditions as EC
+from weakref import ReferenceType, ref
 
 from . import config, cv2find
 from .browsers import BROWSER_CONFIGS, Browser, PageLoadStrategy
@@ -39,6 +43,14 @@
 logger = logging.getLogger(__name__)
 
 
+def _cleanup(bot: ReferenceType[WebBot]):
+    if bot() is not None:
+        try:
+            bot().stop_browser()
+        except Exception:
+            pass
+
+
 class WebBot(BaseBot):
     KEYS = Keys
     DEFAULT_DIMENSIONS = (1600, 900)
@@ -77,6 +89,8 @@ def __init__(self, headless=False):
 
         self._download_folder_path = os.getcwd()
 
+        atexit.register(_cleanup, ref(self))
+
     def __enter__(self):
         pass
 

conftest.py

@@ -7,6 +7,7 @@
 import typing
 import platform
 
+from webdriver_manager.core.driver_cache import DriverCacheManager
 from botcity.web import WebBot, Browser, By, browsers
 from webdriver_manager.chrome import ChromeDriverManager
 from webdriver_manager.firefox import GeckoDriverManager
@@ -97,7 +98,8 @@ def download_driver(request):
     folder_driver = tempfile.mkdtemp()
     browser = request.config.getoption("--browser") or Browser.CHROME
     manager = factory_driver_manager(browser=browser)
-    installed_driver = manager(path=folder_driver).install()
+    cache_manager = DriverCacheManager(root_dir=folder_driver)
+    installed_driver = manager(cache_manager=cache_manager).install()
     yield installed_driver
     shutil.rmtree(folder_driver)