scripts optimization

core_driver/driver.py

@@ -13,14 +13,15 @@
 log = Logger(log_lvl=LogLevel.INFO).get_instance()
 
 
-def _get_driver_path(driver_type=None):
-    if driver_type is None:
+def _get_project_dir():
+    return os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+
+def _get_driver_path(driver_type):
+    if not driver_type:
         ErrorHandler.raise_error(ErrorType.UNSUPPORTED_DRIVER_TYPE)
 
-    project_dir = os.path.dirname(
-        os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
-    )
-    driver_path = os.path.join(project_dir, "resources", driver_type)
+    driver_path = os.path.join(_get_project_dir(), "resources", driver_type)
 
     if not os.path.exists(driver_path):
         ErrorHandler.raise_error(ErrorType.DRIVER_NOT_FOUND, driver_type)
@@ -43,30 +44,31 @@ def create_driver(self, environment, dr_type):
         pass
 
     def get_desired_caps(self, browser="chrome"):
-        caps = YAMLReader.read_caps(browser)
-        log.info(f"Capabilities for driver {caps}")
-        return caps
+        try:
+            caps = YAMLReader.read_caps(browser)
+            log.info(f"Capabilities for {browser} driver: {caps}")
+            return caps
+        except Exception as e:
+            log.error(f"Error reading capabilities for {browser}: {e}")
+            ErrorHandler.raise_error(ErrorType.CAPABILITY_NOT_FOUND, browser)
 
 
 class LocalDriver(Driver):
     def create_driver(self, environment=None, dr_type="chromedriver"):
-        """Tries to use ChromeDriverManager to install the latest driver,
-        and if it fails, it falls back to a locally stored driver in resources."""
         driver = None
+        options = _init_driver_options(dr_type=dr_type)
         try:
             driver_path = ChromeDriverManager().install()
-            options = _init_driver_options(dr_type=dr_type)
             driver = webdriver.Chrome(
-                service=ChromeService(executable_path=driver_path), options=options
-            )
-            log.info(
-                f"Created local Chrome driver with session: {driver.session_id}"
+                service=ChromeService(executable_path=driver_path),
+                options=options
             )
+            log.info(f"Local Chrome driver created with session: {driver.session_id}")
         except Exception as e:
-            log.error(f"Failed to create Chrome driver  {e}")
+            log.error(f"ChromeDriverManager failed, falling back to local driver: {e}")
             driver = webdriver.Chrome(
                 service=ChromeService(_get_driver_path(dr_type)),
-                options=_init_driver_options(dr_type=dr_type),
+                options=options
             )
         _configure_driver(driver, environment)
         return driver
@@ -84,17 +86,17 @@ def create_driver(self, environment=None, dr_type=None):
 
 
 class FirefoxDriver(Driver):
-    def create_driver(self, environment=None, dr_type=None):
+    def create_driver(self, environment=None, dr_type="firefox"):
+        driver = None
+        options = _init_driver_options(dr_type=dr_type)
         try:
-            driver = webdriver.Firefox(options=_init_driver_options(dr_type=dr_type))
-            log.info(f"Created Firefox driver with session: {driver.session_id}")
+            driver = webdriver.Firefox(options=options)
+            log.info(f"Firefox driver created with session: {driver.session_id}")
         except Exception as e:
+            log.error(f"Failed to create Firefox driver, falling back to Chrome: {e}")
             driver = webdriver.Chrome(
-                service=ChromeService(_get_driver_path(dr_type)),
-                options=_init_driver_options(dr_type=dr_type),
-            )
-            log.error(
-                f"Failed to create Firefox driver, falling back to Chrome: {e}"
+                service=ChromeService(_get_driver_path("chromedriver")),
+                options=options
             )
         _configure_driver(driver, environment)
         return driver

pyproject.toml

@@ -6,8 +6,9 @@ authors = ["Dmytro Berezovskyi"]
 readme = "README.md"
 packages = [
     { include = "utils" },
-    { include = "scraper" },
-    { include = "core_driver" }
+    { include = "core_driver" },
+    { include = "scraper", from = "utils" },
+    { include = "cli.py", from="utils"},
 ]
 [tool.poetry.dependencies]
 python = "^3.12"
@@ -26,7 +27,7 @@ setuptools="70.0.0"
 ruff="0.6.8"
 secure-test-automation="^1.3.1"
 colorama="==0.4.6"
-typer="==0.15.1"
+click="==8.1.8"
 
 
 [tool.pytest.ini_options]
@@ -61,3 +62,5 @@ filterwarnings = [
 
 [tool.pytest.config]
 type = ["local", "firefox", "remote"]  # Define a custom command line option for driver types as a list
+[tool.poetry.scripts]
+psaf = "utils.cli:main"

utils/cli.py

@@ -0,0 +1,61 @@
+import argparse
+from pathlib import Path
+
+from utils.scraper.chrome_scraper import ChromePageScraper
+
+
+def main():
+    parser = argparse.ArgumentParser(prog="psaf", description="PSAF CLI Tool")
+
+    # Subcommand for 'get' operations
+    subparsers = parser.add_subparsers(dest="command")
+
+    # Subcommand for 'get chromedriver'
+    get_parser = subparsers.add_parser("get", help="Get a specific resource")
+    get_subparsers = get_parser.add_subparsers(dest="resource")
+
+    # Subcommand for 'chromedriver' under 'get'
+    chromedriver_parser = get_subparsers.add_parser(
+        "chromedriver", help="Download the ChromeDriver"
+    )
+    chromedriver_parser.add_argument(
+        "--platform",
+        type=str,
+        help="Specify the platform (e.g., 'win32', 'linux64').",
+    )
+    chromedriver_parser.add_argument(
+        "--version",
+        type=str,
+        help="Specify the version of ChromeDriver.",
+    )
+    chromedriver_parser.add_argument(
+        "--milestone",
+        type=str,
+        help="Specify the milestone of ChromeDriver.",
+    )
+    chromedriver_parser.add_argument(
+        "--output-dir",
+        type=str,
+        help="Directory to save the downloaded ChromeDriver.",
+    )
+    chromedriver_parser.add_argument(
+        "--extract",
+        action="store_true",
+        help="Extract the downloaded ChromeDriver.",
+    )
+
+    args = parser.parse_args()
+
+    if args.command == "get" and args.resource == "chromedriver":
+        output_dir = Path(args.output_dir) if args.output_dir else None
+        ChromePageScraper.get_chromedriver(
+            platform=args.platform,
+            version=args.version,
+            milestone=args.milestone,
+            d_dir=output_dir,
+            is_extracted=args.extract,
+        )
+
+
+if __name__ == "__main__":
+    main()

utils/error_handler.py

@@ -6,6 +6,7 @@ class ErrorType(Enum):
     EMPTY_URL_ERROR = 2
     UNSUPPORTED_DRIVER_TYPE = 3
     DRIVER_NOT_FOUND = 4
+    CAPABILITY_NOT_FOUND = 5
 
 
 class ErrorHandler:
@@ -14,6 +15,7 @@ class ErrorHandler:
         ErrorType.EMPTY_URL_ERROR: "Environment variable is empty or not found",
         ErrorType.UNSUPPORTED_DRIVER_TYPE: "Unsupported driver type",
         ErrorType.DRIVER_NOT_FOUND: "WebDriver binary not found at ",
+        ErrorType.CAPABILITY_NOT_FOUND: "Capabilities file not found"
     }
 
     @staticmethod

utils/logger.py

@@ -3,6 +3,7 @@
 import time
 from enum import Enum
 from typing import Optional, Callable, Any, Literal
+from threading import Lock
 
 
 class LogLevel(Enum):
@@ -14,23 +15,23 @@ class LogLevel(Enum):
 
 class Singleton(type):
     _instances = {}
+    _lock = Lock()  # Ensure thread safety
 
     def __call__(cls, *args, **kwargs) -> Any:
-        if cls not in cls._instances:
-            cls._instances[cls] = super().__call__(*args, **kwargs)
+        with cls._lock:
+            if cls not in cls._instances:
+                cls._instances[cls] = super().__call__(*args, **kwargs)
         return cls._instances[cls]
 
 
 class Logger(metaclass=Singleton):
-    def __init__(self,
-                 log_lvl: LogLevel = LogLevel.INFO,
-                 log_target: Literal["console", "file", "both"] = "console") -> None:
-        """Logger
-        :param log_lvl: LogLevel
-        :param log_target: Target for logs (store to file, display to console or both(file and console))
-        """
+    def __init__(
+        self,
+        log_lvl: LogLevel = LogLevel.INFO,
+        log_target: Literal["console", "file", "both"] = "console",
+    ) -> None:
         self._log = logging.getLogger("selenium")
-        self._log.setLevel(LogLevel.INFO.value)
+        self._log.setLevel(log_lvl.value)
         self.log_file = self._create_log_file()
         self.log_target = log_target
         self._initialize_logging(log_lvl)
@@ -42,9 +43,7 @@ def _create_log_file(self) -> str:
         )
 
         try:
-            os.makedirs(
-                log_directory, exist_ok=True
-            )  # Create directory if it doesn't exist
+            os.makedirs(log_directory, exist_ok=True)
         except Exception as e:
             raise RuntimeError(
                 f"Failed to create log directory '{log_directory}': {e}"
@@ -56,14 +55,12 @@ def _initialize_logging(self, log_lvl: LogLevel) -> None:
         formatter = logging.Formatter(
             "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
         )
-        # Lot to file or both file and console
         if self.log_target in ("file", "both"):
-            fh = logging.FileHandler(self.log_file, mode="w")
+            fh = logging.FileHandler(self.log_file, mode="a")
             fh.setFormatter(formatter)
             fh.setLevel(log_lvl.value)
             self._log.addHandler(fh)
 
-        # Log to console or both file and console
         if self.log_target in ("console", "both"):
             ch = logging.StreamHandler()
             ch.setFormatter(formatter)
@@ -73,61 +70,30 @@ def _initialize_logging(self, log_lvl: LogLevel) -> None:
     def get_instance(self) -> logging.Logger:
         return self._log
 
-    def annotate(
-        self, message: str, level: Literal["info", "warn", "debug", "error"]
-    ) -> None:
-        """Log a message at the specified level."""
-        if level == "info":
+    def annotate(self, message: str, level: LogLevel) -> None:
+        if level == LogLevel.INFO:
             self._log.info(message)
-        elif level == "warn":
+        elif level == LogLevel.WARNING:
             self._log.warning(message)
-        elif level == "debug":
+        elif level == LogLevel.DEBUG:
             self._log.debug(message)
-        elif level == "error":
+        elif level == LogLevel.ERROR:
             self._log.error(message)
         else:
             raise ValueError(f"Invalid log level: {level}")
 
 
-def log(
-    data: Optional[str] = None,
-    level: Literal["info", "warn", "debug", "error"] = "info",
-) -> Callable:
-    """Decorator to log the current method's execution.
-
-    :param data: Custom log message to use if no docstring is provided.
-    :param level: Level of the logs, e.g., info, warn, debug, error.
-    """
+def log(data: Optional[str] = None, level: LogLevel = LogLevel.INFO) -> Callable:
     logger_instance = Logger()  # Get the singleton instance of Logger
 
     def decorator(func: Callable) -> Callable:
         def wrapper(self, *args, **kwargs) -> Any:
-            # Get the method's docstring
             method_docs = format_method_doc_str(func.__doc__)
-
-            if method_docs is None and data is None:
-                raise ValueError(
-                    f"No documentation available for :: {func.__name__}"
-                )
-
-            # Construct the parameter string for logging
-            params_str = ", ".join(repr(arg) for arg in args)
-            kwargs_str = ", ".join(f"{k}={v!r}" for k, v in kwargs.items())
-            all_params_str = ", ".join(filter(None, [params_str, kwargs_str]))
-
-            logs = (
-                method_docs
-                + f" Method :: {func.__name__}()"
-                + f" with parameters: {all_params_str}"
-                if method_docs
-                else data
-                + f" Method :: {func.__name__}()"
-                + f" with parameters: {all_params_str}"
+            log_message = (
+                (method_docs or data or "")
+                + f" Method :: {func.__name__}() with parameters: {', '.join(map(str, args))}"
             )
-
-            logger_instance.annotate(logs, level)
-
-            # Call the original method, passing *args and **kwargs
+            logger_instance.annotate(log_message, level)
             return func(self, *args, **kwargs)
 
         return wrapper
@@ -136,7 +102,6 @@ def wrapper(self, *args, **kwargs) -> Any:
 
 
 def format_method_doc_str(doc_str: Optional[str]) -> Optional[str]:
-    """Add a dot to the docs string if it doesn't exist."""
     if doc_str and not doc_str.endswith("."):
         return doc_str + "."
     return doc_str