feat: allow to provide the driver and options to be used by Selenium (#316)

This commit is contained in:
Lucas Gomide
2025-05-30 09:13:11 -03:00
committed by GitHub
parent 5d5377cfb9
commit 72b3a8c70a
2 changed files with 29 additions and 8 deletions

View File

@@ -91,9 +91,16 @@ class SeleniumScrapingTool(BaseTool):
"`selenium` and `webdriver-manager` package not found, please run `uv add selenium webdriver-manager`" "`selenium` and `webdriver-manager` package not found, please run `uv add selenium webdriver-manager`"
) )
options: Options = Options() if 'driver' not in kwargs:
options.add_argument("--headless") if 'options' not in kwargs:
self.driver = webdriver.Chrome(options=options) options: Options = Options()
options.add_argument("--headless")
else:
options = kwargs['options']
self.driver = webdriver.Chrome(options=options)
else:
self.driver = kwargs['driver']
self._by = By self._by = By
if cookie is not None: if cookie is not None:
self.cookie = cookie self.cookie = cookie

View File

@@ -2,9 +2,8 @@ import os
import tempfile import tempfile
from unittest.mock import MagicMock, patch from unittest.mock import MagicMock, patch
import pytest
from bs4 import BeautifulSoup from bs4 import BeautifulSoup
from selenium.webdriver.chrome.options import Options
from crewai_tools.tools.selenium_scraping_tool.selenium_scraping_tool import ( from crewai_tools.tools.selenium_scraping_tool.selenium_scraping_tool import (
SeleniumScrapingTool, SeleniumScrapingTool,
) )
@@ -24,9 +23,7 @@ def mock_driver_with_html(html_content):
def initialize_tool_with(mock_driver): def initialize_tool_with(mock_driver):
tool = SeleniumScrapingTool() tool = SeleniumScrapingTool(driver=mock_driver)
tool.driver = mock_driver
return tool return tool
@@ -48,6 +45,17 @@ def test_tool_initialization(mocked_chrome):
except: except:
pass pass
@patch("selenium.webdriver.Chrome")
def test_tool_initialization_with_options(mocked_chrome):
mocked_chrome.return_value = MagicMock()
options = Options()
options.add_argument("--disable-gpu")
SeleniumScrapingTool(options=options)
mocked_chrome.assert_called_once_with(options=options)
@patch("selenium.webdriver.Chrome") @patch("selenium.webdriver.Chrome")
def test_scrape_without_css_selector(_mocked_chrome_driver): def test_scrape_without_css_selector(_mocked_chrome_driver):
@@ -113,3 +121,9 @@ def test_scrape_with_driver_error(_mocked_chrome_driver):
result = tool._run(website_url="https://example.com") result = tool._run(website_url="https://example.com")
assert result == "Error scraping website: WebDriver error occurred" assert result == "Error scraping website: WebDriver error occurred"
mock_driver.close.assert_called_once() mock_driver.close.assert_called_once()
@patch("selenium.webdriver.Chrome")
def test_initialization_with_driver(_mocked_chrome_driver):
mock_driver = MagicMock()
tool = initialize_tool_with(mock_driver)
assert tool.driver == mock_driver