mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-28 17:48:13 +00:00
feat: allow to provide the driver and options to be used by Selenium (#316)
This commit is contained in:
@@ -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
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
Reference in New Issue
Block a user