arcade-mcp/toolkits/firecrawl/tests/test_firecrawl.py

from unittest.mock import patch

import pytest
from arcade_tdk import ToolContext, ToolSecretItem
from arcade_tdk.errors import ToolExecutionError

from arcade_firecrawl.tools import (
    cancel_crawl,
    crawl_website,
    get_crawl_data,
    get_crawl_status,
    map_website,
    scrape_url,
)


@pytest.fixture
def mock_context():
    return ToolContext(secrets=[ToolSecretItem(key="firecrawl_api_key", value="fake_api_key")])


@pytest.fixture
def mock_firecrawl_app_for_scrape():
    with patch("arcade_firecrawl.tools.scrape.FirecrawlApp") as app:
        yield app.return_value


@pytest.fixture
def mock_firecrawl_app_for_crawl():
    with patch("arcade_firecrawl.tools.crawl.FirecrawlApp") as app:
        yield app.return_value


@pytest.fixture
def mock_firecrawl_app_for_map():
    with patch("arcade_firecrawl.tools.map.FirecrawlApp") as app:
        yield app.return_value


@pytest.mark.asyncio
async def test_scrape_url_success(mock_firecrawl_app_for_scrape, mock_context):
    expected_response = {
        "success": True,
        "data": {"scraped_content": "scraped content"},
    }
    mock_firecrawl_app_for_scrape.scrape_url.return_value = expected_response

    result = await scrape_url(mock_context, "http://example.com")
    assert result == expected_response


@pytest.mark.asyncio
async def test_crawl_website_success(mock_firecrawl_app_for_crawl, mock_context):
    expected_response = {
        "id": "12345",
        "success": True,
    }
    mock_firecrawl_app_for_crawl.async_crawl_url.return_value = expected_response
    mock_firecrawl_app_for_crawl.check_crawl_status.return_value = expected_response

    result = await crawl_website(mock_context, "http://example.com")
    assert result == expected_response


@pytest.mark.asyncio
async def test_get_crawl_status_success(mock_firecrawl_app_for_crawl, mock_context):
    expected_response = {"status": "completed"}
    mock_firecrawl_app_for_crawl.check_crawl_status.return_value = expected_response

    result = await get_crawl_status(mock_context, "12345")
    assert result == expected_response


@pytest.mark.asyncio
async def test_get_crawl_data_success(mock_firecrawl_app_for_crawl, mock_context):
    expected_response = {"data": "crawl data"}
    mock_firecrawl_app_for_crawl.check_crawl_status.return_value = expected_response

    result = await get_crawl_data(mock_context, "12345")
    assert result == expected_response


@pytest.mark.asyncio
async def test_cancel_crawl_success(mock_firecrawl_app_for_crawl, mock_context):
    expected_response = {"status": "cancelled"}
    mock_firecrawl_app_for_crawl.cancel_crawl.return_value = expected_response

    result = await cancel_crawl(mock_context, "12345")
    assert result == expected_response


@pytest.mark.asyncio
async def test_map_website_success(mock_firecrawl_app_for_map, mock_context):
    expected_response = {"map": "website map"}
    mock_firecrawl_app_for_map.map_url.return_value = expected_response

    result = await map_website(mock_context, "http://example.com")
    assert result == expected_response


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "method,params,error_message",
    [
        (scrape_url, ("http://example.com",), "Error scraping URL"),
        (crawl_website, ("http://example.com",), "Error crawling website"),
        (get_crawl_status, ("12345",), "Error getting crawl status"),
        (get_crawl_data, ("12345",), "Error getting crawl data"),
        (cancel_crawl, ("12345",), "Error cancelling crawl"),
        (map_website, ("http://example.com",), "Error mapping website"),
    ],
)
async def test_firecrawl_error(
    mock_firecrawl_app_for_scrape,
    mock_firecrawl_app_for_crawl,
    mock_firecrawl_app_for_map,
    mock_context,
    method,
    params,
    error_message,
):
    mock_firecrawl_app_for_scrape.scrape_url.side_effect = Exception(error_message)
    mock_firecrawl_app_for_crawl.async_crawl_url.side_effect = Exception(error_message)
    mock_firecrawl_app_for_crawl.check_crawl_status.side_effect = Exception(error_message)
    mock_firecrawl_app_for_crawl.cancel_crawl.side_effect = Exception(error_message)
    mock_firecrawl_app_for_map.map_url.side_effect = Exception(error_message)

    with pytest.raises(ToolExecutionError):
        await method(mock_context, *params)