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`"
)
options: Options = Options()
options.add_argument("--headless")
self.driver = webdriver.Chrome(options=options)
if 'driver' not in kwargs:
if 'options' not in kwargs:
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
if cookie is not None:
self.cookie = cookie

View File

@@ -2,9 +2,8 @@ import os
import tempfile
from unittest.mock import MagicMock, patch
import pytest
from bs4 import BeautifulSoup
from selenium.webdriver.chrome.options import Options
from crewai_tools.tools.selenium_scraping_tool.selenium_scraping_tool import (
SeleniumScrapingTool,
)
@@ -24,9 +23,7 @@ def mock_driver_with_html(html_content):
def initialize_tool_with(mock_driver):
tool = SeleniumScrapingTool()
tool.driver = mock_driver
tool = SeleniumScrapingTool(driver=mock_driver)
return tool
@@ -48,6 +45,17 @@ def test_tool_initialization(mocked_chrome):
except:
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")
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")
assert result == "Error scraping website: WebDriver error occurred"
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