75fc298681
Updating BrightData - Updating project.toml - Fix linting issues (related to the repo configs)s - Rename package from brightdata -> arcade-brightdata ( also will be used by PyPI) - Added to toolkits.txt so it can be deployed Extra: - Arcade new templates did not have the extra line at the end, so it has been added. --------- Co-authored-by: Francisco Liberal <francisco@arcade.dev>
414 lines
16 KiB
Python
414 lines
16 KiB
Python
from os import environ
|
|
from unittest.mock import Mock, patch
|
|
|
|
import pytest
|
|
import requests
|
|
from arcade_tdk import ToolContext, ToolSecretItem
|
|
from arcade_tdk.errors import ToolExecutionError
|
|
|
|
from arcade_brightdata.bright_data_client import BrightDataClient
|
|
from arcade_brightdata.tools.bright_data_tools import (
|
|
DeviceType,
|
|
SourceType,
|
|
scrape_as_markdown,
|
|
search_engine,
|
|
web_data_feed,
|
|
)
|
|
|
|
BRIGHTDATA_API_KEY = environ.get("TEST_BRIGHTDATA_API_KEY") or "api-key"
|
|
BRIGHTDATA_ZONE = environ.get("TEST_BRIGHTDATA_ZONE") or "unblocker"
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_context():
|
|
context = ToolContext()
|
|
context.secrets = []
|
|
context.secrets.append(ToolSecretItem(key="BRIGHTDATA_API_KEY", value=BRIGHTDATA_API_KEY))
|
|
context.secrets.append(ToolSecretItem(key="BRIGHTDATA_ZONE", value=BRIGHTDATA_ZONE))
|
|
return context
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def cleanup_engines():
|
|
"""Clean up bright data clients after each test to prevent connection leaks."""
|
|
yield
|
|
BrightDataClient.clear_cache()
|
|
|
|
|
|
class TestBrightDataClient:
|
|
def test_get_instance_creates_new_client(self):
|
|
client1 = BrightDataClient.create_client("test_key_1", "zone1")
|
|
client2 = BrightDataClient.create_client("test_key_2", "zone2")
|
|
|
|
assert client1 != client2
|
|
assert client1.api_key == "test_key_1"
|
|
assert client1.zone == "zone1"
|
|
assert client2.api_key == "test_key_2"
|
|
assert client2.zone == "zone2"
|
|
|
|
def test_get_instance_returns_cached_client(self):
|
|
client1 = BrightDataClient.create_client("test_key", "zone1")
|
|
client2 = BrightDataClient.create_client("test_key", "zone1")
|
|
|
|
assert client1 is client2
|
|
|
|
def test_clear_cache(self):
|
|
client1 = BrightDataClient.create_client("test_key", "zone1")
|
|
BrightDataClient.clear_cache()
|
|
client2 = BrightDataClient.create_client("test_key", "zone1")
|
|
|
|
assert client1 is not client2
|
|
|
|
def test_encode_query(self):
|
|
result = BrightDataClient.encode_query("hello world test")
|
|
assert result == "hello%20world%20test"
|
|
|
|
@patch("requests.post")
|
|
def test_make_request_success(self, mock_post):
|
|
mock_response = Mock()
|
|
mock_response.status_code = 200
|
|
mock_response.text = "Success response"
|
|
mock_post.return_value = mock_response
|
|
|
|
client = BrightDataClient("test_key", "test_zone")
|
|
result = client.make_request({"url": "https://example.com"})
|
|
|
|
assert result == "Success response"
|
|
mock_post.assert_called_once()
|
|
|
|
@patch("requests.post")
|
|
def test_make_request_failure(self, mock_post):
|
|
mock_response = Mock()
|
|
mock_response.status_code = 400
|
|
mock_response.text = "Bad Request"
|
|
mock_response.raise_for_status.side_effect = requests.exceptions.HTTPError(
|
|
"400 Client Error"
|
|
)
|
|
mock_post.return_value = mock_response
|
|
|
|
client = BrightDataClient("test_key", "test_zone")
|
|
|
|
with pytest.raises(requests.exceptions.HTTPError):
|
|
client.make_request({"url": "https://example.com"})
|
|
|
|
|
|
class TestScrapeAsMarkdown:
|
|
@patch("arcade_brightdata.tools.bright_data_tools.BrightDataClient")
|
|
def test_scrape_as_markdown_success(self, mock_engine_class, mock_context):
|
|
mock_client = Mock()
|
|
mock_client.make_request.return_value = "# Test Page\n\nContent here"
|
|
mock_engine_class.create_client.return_value = mock_client
|
|
|
|
result = scrape_as_markdown(mock_context, "https://example.com")
|
|
|
|
assert result == "# Test Page\n\nContent here"
|
|
mock_engine_class.create_client.assert_called_once_with(
|
|
api_key=BRIGHTDATA_API_KEY, zone=BRIGHTDATA_ZONE
|
|
)
|
|
mock_client.make_request.assert_called_once_with({
|
|
"url": "https://example.com",
|
|
"zone": BRIGHTDATA_ZONE,
|
|
"format": "raw",
|
|
"data_format": "markdown",
|
|
})
|
|
|
|
|
|
class TestSearchEngine:
|
|
@patch("arcade_brightdata.tools.bright_data_tools.BrightDataClient")
|
|
def test_search_engine_google_basic(self, mock_engine_class, mock_context):
|
|
mock_client = Mock()
|
|
mock_client.make_request.return_value = "# Search Results\n\nResult 1\nResult 2"
|
|
mock_engine_class.create_client.return_value = mock_client
|
|
mock_engine_class.encode_query.return_value = "test%20query"
|
|
|
|
result = search_engine(mock_context, "test query")
|
|
|
|
assert result == "# Search Results\n\nResult 1\nResult 2"
|
|
mock_engine_class.create_client.assert_called_once_with(
|
|
api_key=BRIGHTDATA_API_KEY, zone=BRIGHTDATA_ZONE
|
|
)
|
|
|
|
@patch("arcade_brightdata.tools.bright_data_tools.BrightDataClient")
|
|
def test_search_engine_bing(self, mock_engine_class, mock_context):
|
|
mock_client = Mock()
|
|
mock_client.make_request.return_value = "# Bing Results"
|
|
mock_engine_class.create_client.return_value = mock_client
|
|
mock_engine_class.encode_query.return_value = "test%20query"
|
|
|
|
result = search_engine(mock_context, "test query", engine="bing")
|
|
|
|
assert result == "# Bing Results"
|
|
expected_payload = {
|
|
"url": "https://www.bing.com/search?q=test%20query",
|
|
"zone": BRIGHTDATA_ZONE,
|
|
"format": "raw",
|
|
"data_format": "markdown",
|
|
}
|
|
mock_client.make_request.assert_called_once_with(expected_payload)
|
|
|
|
@patch("arcade_brightdata.tools.bright_data_tools.BrightDataClient")
|
|
def test_search_engine_google_with_parameters(self, mock_engine_class, mock_context):
|
|
mock_client = Mock()
|
|
mock_client.make_request.return_value = "# Google Results with params"
|
|
mock_engine_class.create_client.return_value = mock_client
|
|
mock_engine_class.encode_query.side_effect = lambda x: x.replace(" ", "%20")
|
|
|
|
result = search_engine(
|
|
mock_context,
|
|
"test query",
|
|
language="en",
|
|
country_code="us",
|
|
search_type="images",
|
|
start=10,
|
|
num_results=20,
|
|
location="New York",
|
|
device=DeviceType.MOBILE,
|
|
return_json=True,
|
|
)
|
|
|
|
assert result == "# Google Results with params"
|
|
call_args = mock_client.make_request.call_args[0][0]
|
|
|
|
assert "hl=en" in call_args["url"]
|
|
assert "gl=us" in call_args["url"]
|
|
assert "tbm=isch" in call_args["url"]
|
|
assert "start=10" in call_args["url"]
|
|
assert "num=20" in call_args["url"]
|
|
assert "brd_mobile=1" in call_args["url"]
|
|
assert "brd_json=1" in call_args["url"]
|
|
assert call_args["data_format"] == "raw"
|
|
|
|
def test_search_engine_invalid_engine(self, mock_context):
|
|
with pytest.raises(ToolExecutionError):
|
|
search_engine(mock_context, "test query", engine="invalid_engine")
|
|
|
|
@patch("arcade_brightdata.tools.bright_data_tools.BrightDataClient")
|
|
def test_search_engine_google_jobs(self, mock_engine_class, mock_context):
|
|
mock_client = Mock()
|
|
mock_client.make_request.return_value = "# Job Results"
|
|
mock_engine_class.create_client.return_value = mock_client
|
|
mock_engine_class.encode_query.return_value = "python%20developer"
|
|
|
|
result = search_engine(mock_context, "python developer", search_type="jobs")
|
|
|
|
assert result == "# Job Results"
|
|
call_args = mock_client.make_request.call_args[0][0]
|
|
assert "ibp=htl;jobs" in call_args["url"]
|
|
|
|
|
|
class TestWebDataFeed:
|
|
@patch("arcade_brightdata.tools.bright_data_tools._extract_structured_data")
|
|
@patch("arcade_brightdata.tools.bright_data_tools.BrightDataClient")
|
|
def test_web_data_feed_success(self, mock_engine_class, mock_extract, mock_context):
|
|
mock_client = Mock()
|
|
mock_engine_class.create_client.return_value = mock_client
|
|
mock_extract.return_value = {"title": "Test Product", "price": "$19.99"}
|
|
|
|
result = web_data_feed(mock_context, "amazon_product", "https://amazon.com/dp/B08N5WRWNW")
|
|
|
|
expected_json = '{\n "title": "Test Product",\n "price": "$19.99"\n}'
|
|
assert result == expected_json
|
|
|
|
mock_engine_class.create_client.assert_called_once_with(api_key=BRIGHTDATA_API_KEY)
|
|
mock_extract.assert_called_once_with(
|
|
client=mock_client,
|
|
source_type=SourceType.AMAZON_PRODUCT,
|
|
url="https://amazon.com/dp/B08N5WRWNW",
|
|
num_of_reviews=None,
|
|
timeout=600,
|
|
polling_interval=1,
|
|
)
|
|
|
|
@patch("arcade_brightdata.tools.bright_data_tools._extract_structured_data")
|
|
@patch("arcade_brightdata.tools.bright_data_tools.BrightDataClient")
|
|
def test_web_data_feed_with_reviews(self, mock_engine_class, mock_extract, mock_context):
|
|
mock_client = Mock()
|
|
mock_engine_class.create_client.return_value = mock_client
|
|
mock_extract.return_value = [{"review": "Great product!", "rating": 5}]
|
|
|
|
result = web_data_feed(
|
|
mock_context,
|
|
"facebook_company_reviews",
|
|
"https://facebook.com/company",
|
|
num_of_reviews=50,
|
|
timeout=300,
|
|
polling_interval=2,
|
|
)
|
|
|
|
expected_json = '[\n {\n "review": "Great product!",\n "rating": 5\n }\n]'
|
|
assert result == expected_json
|
|
|
|
mock_extract.assert_called_once_with(
|
|
client=mock_client,
|
|
source_type=SourceType.FACEBOOK_COMPANY_REVIEWS,
|
|
url="https://facebook.com/company",
|
|
num_of_reviews=50,
|
|
timeout=300,
|
|
polling_interval=2,
|
|
)
|
|
|
|
|
|
class TestExtractStructuredData:
|
|
@patch("requests.get")
|
|
@patch("requests.post")
|
|
def test_extract_structured_data_success(self, mock_post, mock_get):
|
|
from arcade_brightdata.tools.bright_data_tools import _extract_structured_data
|
|
|
|
client = BrightDataClient("test_key", "test_zone")
|
|
|
|
mock_trigger_response = Mock()
|
|
mock_trigger_response.json.return_value = {"snapshot_id": "snap_123"}
|
|
mock_post.return_value = mock_trigger_response
|
|
|
|
mock_snapshot_response = Mock()
|
|
mock_snapshot_response.json.return_value = {"data": "extracted_data"}
|
|
mock_get.return_value = mock_snapshot_response
|
|
|
|
result = _extract_structured_data(
|
|
client=client,
|
|
source_type=SourceType.AMAZON_PRODUCT,
|
|
url="https://amazon.com/dp/TEST",
|
|
timeout=10,
|
|
polling_interval=0.1,
|
|
)
|
|
|
|
assert result == {"data": "extracted_data"}
|
|
|
|
mock_post.assert_called_once()
|
|
trigger_call = mock_post.call_args
|
|
assert "gd_l7q7dkf244hwjntr0" in str(trigger_call) # Amazon product dataset ID
|
|
|
|
mock_get.assert_called_once()
|
|
snapshot_call = mock_get.call_args
|
|
assert "snap_123" in str(snapshot_call)
|
|
|
|
@patch("requests.get")
|
|
@patch("requests.post")
|
|
def test_extract_structured_data_with_polling(self, mock_post, mock_get):
|
|
from arcade_brightdata.tools.bright_data_tools import _extract_structured_data
|
|
|
|
client = BrightDataClient("test_key", "test_zone")
|
|
|
|
mock_trigger_response = Mock()
|
|
mock_trigger_response.json.return_value = {"snapshot_id": "snap_123"}
|
|
mock_post.return_value = mock_trigger_response
|
|
|
|
running_response = Mock()
|
|
running_response.json.return_value = {"status": "running"}
|
|
|
|
complete_response = Mock()
|
|
complete_response.json.return_value = {"data": "final_data"}
|
|
|
|
mock_get.side_effect = [running_response, complete_response]
|
|
|
|
result = _extract_structured_data(
|
|
client=client,
|
|
source_type=SourceType.LINKEDIN_PERSON_PROFILE,
|
|
url="https://linkedin.com/in/test",
|
|
timeout=10,
|
|
polling_interval=0.1,
|
|
)
|
|
|
|
assert result == {"data": "final_data"}
|
|
assert mock_get.call_count == 2
|
|
|
|
@patch("requests.post")
|
|
def test_extract_structured_data_invalid_source_type(self, mock_post):
|
|
from arcade_brightdata.tools.bright_data_tools import _extract_structured_data
|
|
|
|
client = BrightDataClient("test_key", "test_zone")
|
|
|
|
# Create a mock SourceType that doesn't exist in the datasets dict
|
|
class InvalidSourceType:
|
|
value = "invalid_source"
|
|
|
|
with pytest.raises(KeyError):
|
|
_extract_structured_data(
|
|
client=client, source_type=InvalidSourceType(), url="https://example.com"
|
|
)
|
|
|
|
@patch("requests.get")
|
|
@patch("requests.post")
|
|
def test_extract_structured_data_no_snapshot_id(self, mock_post, mock_get):
|
|
from arcade_brightdata.tools.bright_data_tools import _extract_structured_data
|
|
|
|
client = BrightDataClient("test_key", "test_zone")
|
|
|
|
# Mock trigger response without snapshot_id
|
|
mock_trigger_response = Mock()
|
|
mock_trigger_response.json.return_value = {}
|
|
mock_post.return_value = mock_trigger_response
|
|
|
|
with pytest.raises(Exception) as exc_info:
|
|
_extract_structured_data(
|
|
client=client,
|
|
source_type=SourceType.AMAZON_PRODUCT,
|
|
url="https://amazon.com/dp/TEST",
|
|
)
|
|
|
|
assert "No snapshot ID returned from trigger request" in str(exc_info.value)
|
|
|
|
@patch("requests.get")
|
|
@patch("requests.post")
|
|
@patch("time.sleep")
|
|
def test_extract_structured_data_timeout(self, mock_sleep, mock_post, mock_get):
|
|
from arcade_brightdata.tools.bright_data_tools import _extract_structured_data
|
|
|
|
client = BrightDataClient("test_key", "test_zone")
|
|
|
|
# Mock trigger response
|
|
mock_trigger_response = Mock()
|
|
mock_trigger_response.json.return_value = {"snapshot_id": "snap_123"}
|
|
mock_post.return_value = mock_trigger_response
|
|
|
|
# Mock snapshot response that always returns running
|
|
mock_snapshot_response = Mock()
|
|
mock_snapshot_response.json.return_value = {"status": "running"}
|
|
mock_get.return_value = mock_snapshot_response
|
|
|
|
with pytest.raises(TimeoutError) as exc_info:
|
|
_extract_structured_data(
|
|
client=client,
|
|
source_type=SourceType.AMAZON_PRODUCT,
|
|
url="https://amazon.com/dp/TEST",
|
|
timeout=2,
|
|
polling_interval=0.1,
|
|
)
|
|
|
|
assert "Timeout after 2 seconds waiting for amazon_product data" in str(exc_info.value)
|
|
|
|
|
|
class TestIntegration:
|
|
"""Integration tests that test the full flow without mocking internal components."""
|
|
|
|
@patch("requests.post")
|
|
def test_scrape_as_markdown_integration(self, mock_post, mock_context):
|
|
mock_response = Mock()
|
|
mock_response.status_code = 200
|
|
mock_response.text = "# Integration Test\n\nThis is a test page"
|
|
mock_post.return_value = mock_response
|
|
|
|
result = scrape_as_markdown(mock_context, "https://example.com")
|
|
|
|
assert result == "# Integration Test\n\nThis is a test page"
|
|
|
|
# Verify the request was made correctly
|
|
call_args = mock_post.call_args
|
|
assert call_args[1]["headers"]["Authorization"] == f"Bearer {BRIGHTDATA_API_KEY}"
|
|
assert "https://api.brightdata.com/request" in str(call_args)
|
|
|
|
@patch("requests.post")
|
|
def test_search_engine_integration(self, mock_post, mock_context):
|
|
mock_response = Mock()
|
|
mock_response.status_code = 200
|
|
mock_response.text = "# Search Results\n\n1. First result\n2. Second result"
|
|
mock_post.return_value = mock_response
|
|
|
|
result = search_engine(mock_context, "test query", engine="google")
|
|
|
|
assert result == "# Search Results\n\n1. First result\n2. Second result"
|
|
|
|
call_args = mock_post.call_args
|
|
payload = call_args[1]["data"]
|
|
assert '"url": "https://www.google.com/search?q=test%20query' in payload
|
|
assert '"data_format": "markdown"' in payload
|