mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-03 07:23:33 +00:00
test: add unit tests for fetch registry, management, firecrawl extended, callback utils
Add comprehensive test coverage for: - fetch_tool_registry.py (31 tests, CRUD operations) - fetch_tool_management.py (20 tests, API endpoints) - firecrawl/fetch/transformation.py (14 tests, sync/async/self-hosted) - proxy/common_utils/callback_utils.py (4 tests, callback init) All tests use pure mocks (MagicMock/AsyncMock) without real DB, Prisma, or heavy proxy dependencies.
This commit is contained in:
Thomas Menard
parent
4591b52efb
commit
2ad2c85cc3
4 changed files with 908 additions and 0 deletions
117
tests/webfetch_tests/test_callback_utils_webfetch.py
Normal file
117
tests/webfetch_tests/test_callback_utils_webfetch.py
Normal file
|
|
@ -0,0 +1,117 @@
|
|||
"""Test callback_utils webfetch_interception initialization.
|
||||
|
||||
This file tests the 4 lines that register the webfetch_interception callback
|
||||
in the proxy's callback initialization.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
|
||||
class TestInitializeCallbacksWebFetch:
|
||||
"""Test that webfetch_interception callback is initialized correctly."""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def setup(self, monkeypatch):
|
||||
"""Set up mocks for heavy dependencies."""
|
||||
# Mock proxy_server.prisma_client to avoid DB setup
|
||||
self.mock_prisma = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.prisma_client",
|
||||
self.mock_prisma,
|
||||
raising=False,
|
||||
)
|
||||
|
||||
@patch("litellm.integrations.webfetch_interception.handler.WebFetchInterceptionLogger.initialize_from_proxy_config")
|
||||
def test_webfetch_interception_callback_registered(self, mock_init):
|
||||
"""Test that webfetch_interception callback creates and appends logger."""
|
||||
from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy
|
||||
|
||||
mock_logger = MagicMock()
|
||||
mock_init.return_value = mock_logger
|
||||
|
||||
litellm_settings = {
|
||||
"webfetch_interception_params": {
|
||||
"fetch_provider": "firecrawl",
|
||||
"api_key": "fc-test",
|
||||
}
|
||||
}
|
||||
|
||||
result = initialize_callbacks_on_proxy(
|
||||
value=["webfetch_interception"],
|
||||
premium_user=False,
|
||||
config_file_path="/dev/null",
|
||||
litellm_settings=litellm_settings,
|
||||
)
|
||||
|
||||
mock_init.assert_called_once_with(
|
||||
litellm_settings=litellm_settings,
|
||||
callback_specific_params={},
|
||||
)
|
||||
assert mock_logger in result
|
||||
|
||||
@patch("litellm.integrations.webfetch_interception.handler.WebFetchInterceptionLogger.initialize_from_proxy_config")
|
||||
def test_webfetch_interception_with_callback_specific_params(self, mock_init):
|
||||
"""Test with callback_specific_params containing webfetch_interception config."""
|
||||
from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy
|
||||
|
||||
mock_logger = MagicMock()
|
||||
mock_init.return_value = mock_logger
|
||||
|
||||
callback_specific_params = {
|
||||
"webfetch_interception": {
|
||||
"fetch_provider": "firecrawl",
|
||||
"api_key": "fc-specific",
|
||||
}
|
||||
}
|
||||
|
||||
result = initialize_callbacks_on_proxy(
|
||||
value=["webfetch_interception"],
|
||||
premium_user=False,
|
||||
config_file_path="/dev/null",
|
||||
litellm_settings={},
|
||||
callback_specific_params=callback_specific_params,
|
||||
)
|
||||
|
||||
mock_init.assert_called_once_with(
|
||||
litellm_settings={},
|
||||
callback_specific_params=callback_specific_params,
|
||||
)
|
||||
assert mock_logger in result
|
||||
|
||||
@patch("litellm.integrations.webfetch_interception.handler.WebFetchInterceptionLogger.initialize_from_proxy_config")
|
||||
def test_webfetch_interception_not_in_list(self, mock_init):
|
||||
"""Test that other callbacks don't trigger webfetch initialization."""
|
||||
from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy
|
||||
|
||||
mock_logger = MagicMock()
|
||||
mock_init.return_value = mock_logger
|
||||
|
||||
result = initialize_callbacks_on_proxy(
|
||||
value=["other_callback"],
|
||||
premium_user=False,
|
||||
config_file_path="/dev/null",
|
||||
litellm_settings={},
|
||||
)
|
||||
|
||||
assert mock_logger not in result
|
||||
|
||||
def test_webfetch_interception_empty_params(self):
|
||||
"""Test that empty params still work."""
|
||||
from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy
|
||||
|
||||
with patch(
|
||||
"litellm.integrations.webfetch_interception.handler.WebFetchInterceptionLogger.initialize_from_proxy_config"
|
||||
) as mock_init:
|
||||
mock_logger = MagicMock()
|
||||
mock_init.return_value = mock_logger
|
||||
|
||||
result = initialize_callbacks_on_proxy(
|
||||
value=["webfetch_interception"],
|
||||
premium_user=False,
|
||||
config_file_path="/dev/null",
|
||||
litellm_settings={},
|
||||
)
|
||||
|
||||
mock_init.assert_called_once()
|
||||
assert mock_logger in result
|
||||
256
tests/webfetch_tests/test_fetch_tool_management.py
Normal file
256
tests/webfetch_tests/test_fetch_tool_management.py
Normal file
|
|
@ -0,0 +1,256 @@
|
|||
"""Unit tests for fetch_tool_management.py endpoints.
|
||||
|
||||
All heavy dependencies (FastAPI, proxy_server, auth) are mocked.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import MagicMock, AsyncMock, patch, PropertyMock
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
|
||||
class TestGetEndpoints:
|
||||
"""Test GET endpoints."""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def setup(self, monkeypatch):
|
||||
"""Mock proxy_server.prisma_client before importing."""
|
||||
self.mock_prisma = MagicMock()
|
||||
self.mock_config = MagicMock()
|
||||
self.mock_prisma_client_module = MagicMock()
|
||||
self.mock_prisma_client_module.prisma_client = self.mock_prisma
|
||||
self.mock_prisma_client_module.proxy_config = self.mock_config
|
||||
|
||||
self.mock_registry = MagicMock()
|
||||
self.mock_registry.get_all_fetch_tools_from_db = AsyncMock(return_value=[])
|
||||
self.mock_registry.get_fetch_tool_from_db = AsyncMock(return_value=None)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server",
|
||||
self.mock_prisma_client_module,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.fetch_endpoints.fetch_tool_management.FETCH_TOOL_REGISTRY",
|
||||
self.mock_registry,
|
||||
)
|
||||
|
||||
# Now import after patching
|
||||
from litellm.proxy.fetch_endpoints import fetch_tool_management
|
||||
self.module = fetch_tool_management
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_fetch_tools_success_empty(self, setup):
|
||||
"""Test listing with empty DB."""
|
||||
result = await self.module.list_fetch_tools()
|
||||
assert result == {"fetch_tools": []}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_fetch_tools_with_tools(self, setup):
|
||||
"""Test listing with tools from DB."""
|
||||
mock_tool = {
|
||||
"fetch_tool_id": "1",
|
||||
"fetch_tool_name": "test-tool",
|
||||
"litellm_params": {"api_key": "sk-secret"},
|
||||
"fetch_tool_info": {},
|
||||
}
|
||||
self.mock_registry.get_all_fetch_tools_from_db = AsyncMock(
|
||||
return_value=[mock_tool]
|
||||
)
|
||||
|
||||
result = await self.module.list_fetch_tools()
|
||||
tools = result["fetch_tools"]
|
||||
assert len(tools) == 1
|
||||
assert tools[0]["fetch_tool_name"] == "test-tool"
|
||||
# API keys are masked by _get_masked_values
|
||||
assert "api_key" in tools[0]["litellm_params"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_fetch_tools_with_config_tools(self, setup):
|
||||
"""Test listing merges DB + config tools."""
|
||||
db_tool = {
|
||||
"fetch_tool_id": "1",
|
||||
"fetch_tool_name": "db-tool",
|
||||
"litellm_params": {},
|
||||
"fetch_tool_info": {},
|
||||
}
|
||||
self.mock_registry.get_all_fetch_tools_from_db = AsyncMock(
|
||||
return_value=[db_tool]
|
||||
)
|
||||
|
||||
# Config has a tool with same name → not added twice
|
||||
self.mock_config.config = {"fetch_tools": [{"fetch_tool_name": "db-tool"}]}
|
||||
self.mock_config.parse_fetch_tools = MagicMock(return_value=[])
|
||||
|
||||
result = await self.module.list_fetch_tools()
|
||||
assert len(result["fetch_tools"]) == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_fetch_tool_found(self, setup):
|
||||
"""Test getting a tool that exists."""
|
||||
self.mock_registry.get_fetch_tool_from_db = AsyncMock(
|
||||
return_value={
|
||||
"fetch_tool_id": "123",
|
||||
"fetch_tool_name": "test",
|
||||
}
|
||||
)
|
||||
result = await self.module.get_fetch_tool("123")
|
||||
assert result["fetch_tool_id"] == "123"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_fetch_tool_not_found(self, setup):
|
||||
"""Test getting a tool that does not exist."""
|
||||
self.mock_registry.get_fetch_tool_from_db = AsyncMock(return_value=None)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await self.module.get_fetch_tool("missing")
|
||||
assert exc_info.value.status_code == 404
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_available_providers(self, setup):
|
||||
"""Test listing available providers."""
|
||||
result = await self.module.list_available_fetch_providers()
|
||||
assert "providers" in result
|
||||
assert len(result["providers"]) >= 1
|
||||
assert result["providers"][0]["provider_name"] == "firecrawl"
|
||||
|
||||
|
||||
class TestCreateEndpoint:
|
||||
"""Test POST /fetch_tools."""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def setup(self, monkeypatch):
|
||||
self.mock_prisma = MagicMock()
|
||||
self.mock_prisma_client_module = MagicMock()
|
||||
self.mock_prisma_client_module.prisma_client = self.mock_prisma
|
||||
|
||||
self.mock_registry = MagicMock()
|
||||
self.mock_registry.add_fetch_tool_to_db = AsyncMock(
|
||||
return_value={"fetch_tool_id": "new-id", "fetch_tool_name": "new-tool"}
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server",
|
||||
self.mock_prisma_client_module,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.fetch_endpoints.fetch_tool_management.FETCH_TOOL_REGISTRY",
|
||||
self.mock_registry,
|
||||
)
|
||||
|
||||
from litellm.proxy.fetch_endpoints import fetch_tool_management
|
||||
self.module = fetch_tool_management
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_success(self, setup):
|
||||
"""Test creating a tool."""
|
||||
request = {
|
||||
"fetch_tool": {
|
||||
"fetch_tool_name": "new-tool",
|
||||
"litellm_params": {"provider": "firecrawl"},
|
||||
}
|
||||
}
|
||||
result = await self.module.create_fetch_tool(request)
|
||||
assert result["fetch_tool_id"] == "new-id"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_error(self, setup):
|
||||
"""Test creation error raises HTTPException."""
|
||||
self.mock_registry.add_fetch_tool_to_db = AsyncMock(
|
||||
side_effect=Exception("db fail")
|
||||
)
|
||||
|
||||
request = {"fetch_tool": {"fetch_tool_name": "fail"}}
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await self.module.create_fetch_tool(request)
|
||||
assert exc_info.value.status_code == 500
|
||||
|
||||
|
||||
class TestUpdateEndpoint:
|
||||
"""Test PUT /fetch_tools/{id}."""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def setup(self, monkeypatch):
|
||||
self.mock_prisma = MagicMock()
|
||||
self.mock_prisma_client_module = MagicMock()
|
||||
self.mock_prisma_client_module.prisma_client = self.mock_prisma
|
||||
|
||||
self.mock_registry = MagicMock()
|
||||
self.mock_registry.update_fetch_tool_in_db = AsyncMock(
|
||||
return_value={"fetch_tool_id": "123", "fetch_tool_name": "updated"}
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server",
|
||||
self.mock_prisma_client_module,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.fetch_endpoints.fetch_tool_management.FETCH_TOOL_REGISTRY",
|
||||
self.mock_registry,
|
||||
)
|
||||
|
||||
from litellm.proxy.fetch_endpoints import fetch_tool_management
|
||||
self.module = fetch_tool_management
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_success(self, setup):
|
||||
"""Test updating a tool."""
|
||||
request = {"fetch_tool": {"fetch_tool_name": "updated", "litellm_params": {}}}
|
||||
result = await self.module.update_fetch_tool("123", request)
|
||||
assert result["fetch_tool_name"] == "updated"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_error(self, setup):
|
||||
"""Test update error."""
|
||||
self.mock_registry.update_fetch_tool_in_db = AsyncMock(
|
||||
side_effect=Exception("not found")
|
||||
)
|
||||
|
||||
request = {"fetch_tool": {"fetch_tool_name": "updated"}}
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await self.module.update_fetch_tool("bad-id", request)
|
||||
assert exc_info.value.status_code == 500
|
||||
|
||||
|
||||
class TestDeleteEndpoint:
|
||||
"""Test DELETE /fetch_tools/{id}."""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def setup(self, monkeypatch):
|
||||
self.mock_prisma = MagicMock()
|
||||
self.mock_prisma_client_module = MagicMock()
|
||||
self.mock_prisma_client_module.prisma_client = self.mock_prisma
|
||||
|
||||
self.mock_registry = MagicMock()
|
||||
self.mock_registry.delete_fetch_tool_from_db = AsyncMock(
|
||||
return_value={"message": "deleted"}
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server",
|
||||
self.mock_prisma_client_module,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.fetch_endpoints.fetch_tool_management.FETCH_TOOL_REGISTRY",
|
||||
self.mock_registry,
|
||||
)
|
||||
|
||||
from litellm.proxy.fetch_endpoints import fetch_tool_management
|
||||
self.module = fetch_tool_management
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_success(self, setup):
|
||||
"""Test deleting a tool."""
|
||||
result = await self.module.delete_fetch_tool("123")
|
||||
assert result["message"] == "deleted"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_error(self, setup):
|
||||
"""Test delete error."""
|
||||
self.mock_registry.delete_fetch_tool_from_db = AsyncMock(
|
||||
side_effect=Exception("not found")
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await self.module.delete_fetch_tool("bad-id")
|
||||
assert exc_info.value.status_code == 500
|
||||
317
tests/webfetch_tests/test_fetch_tool_registry.py
Normal file
317
tests/webfetch_tests/test_fetch_tool_registry.py
Normal file
|
|
@ -0,0 +1,317 @@
|
|||
"""Unit tests for fetch_tool_registry.
|
||||
|
||||
These tests use pure mocks without real DB access.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import MagicMock, AsyncMock
|
||||
|
||||
from litellm.proxy.fetch_endpoints.fetch_tool_registry import FetchToolRegistry
|
||||
|
||||
|
||||
class TestConvertPrismaToDict:
|
||||
"""Test _convert_prisma_to_dict static method."""
|
||||
|
||||
def test_empty_object(self):
|
||||
"""Test with empty object."""
|
||||
obj = MagicMock()
|
||||
obj.__iter__ = MagicMock(return_value=iter([]))
|
||||
result = FetchToolRegistry._convert_prisma_to_dict(obj)
|
||||
assert result == {}
|
||||
|
||||
def test_with_datetimes(self):
|
||||
"""Test conversion of datetime fields."""
|
||||
from datetime import datetime, timezone
|
||||
|
||||
obj = MagicMock()
|
||||
now = datetime.now(timezone.utc)
|
||||
obj.__dict__ = {
|
||||
"fetch_tool_id": "123",
|
||||
"fetch_tool_name": "test",
|
||||
"created_at": now,
|
||||
"updated_at": now,
|
||||
}
|
||||
obj.__iter__ = MagicMock(return_value=iter(obj.__dict__.items()))
|
||||
|
||||
result = FetchToolRegistry._convert_prisma_to_dict(obj)
|
||||
|
||||
assert result["fetch_tool_id"] == "123"
|
||||
assert result["fetch_tool_name"] == "test"
|
||||
assert result["created_at"] == now.isoformat()
|
||||
assert result["updated_at"] == now.isoformat()
|
||||
|
||||
def test_with_none_datetimes(self):
|
||||
"""Test with None datetime fields."""
|
||||
from datetime import datetime, timezone
|
||||
|
||||
obj = MagicMock()
|
||||
now = datetime.now(timezone.utc)
|
||||
obj.__dict__ = {
|
||||
"created_at": now,
|
||||
"updated_at": None,
|
||||
}
|
||||
obj.__iter__ = MagicMock(return_value=iter(obj.__dict__.items()))
|
||||
|
||||
result = FetchToolRegistry._convert_prisma_to_dict(obj)
|
||||
|
||||
assert result["created_at"] == now.isoformat()
|
||||
assert result["updated_at"] is None
|
||||
|
||||
def test_no_datetime_fields(self):
|
||||
"""Test without any datetime fields."""
|
||||
obj = MagicMock()
|
||||
obj.__dict__ = {"foo": "bar"}
|
||||
obj.__iter__ = MagicMock(return_value=iter(obj.__dict__.items()))
|
||||
|
||||
result = FetchToolRegistry._convert_prisma_to_dict(obj)
|
||||
|
||||
assert result == {"foo": "bar"}
|
||||
|
||||
|
||||
class TestAddFetchTool:
|
||||
"""Test add_fetch_tool_to_db."""
|
||||
|
||||
@pytest.fixture
|
||||
def registry(self):
|
||||
return FetchToolRegistry()
|
||||
|
||||
@pytest.fixture
|
||||
def mock_prisma(self):
|
||||
prisma = MagicMock()
|
||||
prisma.db = MagicMock()
|
||||
prisma.db.litellm_fetchtoolstable = MagicMock()
|
||||
return prisma
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_success(self, registry, mock_prisma):
|
||||
"""Test successful creation."""
|
||||
mock_tool = MagicMock()
|
||||
mock_tool.fetch_tool_id = "abc-123"
|
||||
mock_tool.created_at = MagicMock()
|
||||
mock_tool.created_at.isoformat.return_value = "2026-05-03T12:00:00"
|
||||
mock_tool.updated_at = MagicMock()
|
||||
mock_tool.updated_at.isoformat.return_value = "2026-05-03T12:00:00"
|
||||
|
||||
mock_prisma.db.litellm_fetchtoolstable.create = AsyncMock(return_value=mock_tool)
|
||||
|
||||
fetch_tool = {
|
||||
"fetch_tool_name": "my-tool",
|
||||
"litellm_params": {"provider": "firecrawl"},
|
||||
"fetch_tool_info": {"description": "test"},
|
||||
}
|
||||
|
||||
result = await registry.add_fetch_tool_to_db(fetch_tool, mock_prisma)
|
||||
|
||||
assert result["fetch_tool_id"] == "abc-123"
|
||||
assert result["fetch_tool_name"] == "my-tool"
|
||||
mock_prisma.db.litellm_fetchtoolstable.create.assert_awaited_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_error_handling(self, registry, mock_prisma):
|
||||
"""Test error handling."""
|
||||
mock_prisma.db.litellm_fetchtoolstable.create = AsyncMock(
|
||||
side_effect=Exception("DB error")
|
||||
)
|
||||
|
||||
fetch_tool = {"fetch_tool_name": "fail"}
|
||||
|
||||
with pytest.raises(Exception, match="DB error"):
|
||||
await registry.add_fetch_tool_to_db(fetch_tool, mock_prisma)
|
||||
|
||||
|
||||
class TestDeleteFetchTool:
|
||||
"""Test delete_fetch_tool_from_db."""
|
||||
|
||||
@pytest.fixture
|
||||
def registry(self):
|
||||
return FetchToolRegistry()
|
||||
|
||||
@pytest.fixture
|
||||
def mock_prisma(self):
|
||||
prisma = MagicMock()
|
||||
prisma.db = MagicMock()
|
||||
prisma.db.litellm_fetchtoolstable = MagicMock()
|
||||
return prisma
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_success(self, registry, mock_prisma):
|
||||
"""Test successful deletion."""
|
||||
mock_prisma.db.litellm_fetchtoolstable.delete = AsyncMock(return_value=None)
|
||||
|
||||
result = await registry.delete_fetch_tool_from_db("tool-123", mock_prisma)
|
||||
|
||||
assert result["message"] == "Fetch tool tool-123 deleted successfully"
|
||||
mock_prisma.db.litellm_fetchtoolstable.delete.assert_awaited_once_with(
|
||||
where={"fetch_tool_id": "tool-123"}
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_error(self, registry, mock_prisma):
|
||||
"""Test deletion error."""
|
||||
mock_prisma.db.litellm_fetchtoolstable.delete = AsyncMock(
|
||||
side_effect=Exception("not found")
|
||||
)
|
||||
|
||||
with pytest.raises(Exception, match="not found"):
|
||||
await registry.delete_fetch_tool_from_db("bad-id", mock_prisma)
|
||||
|
||||
|
||||
class TestGetFetchTool:
|
||||
"""Test get_fetch_tool_from_db."""
|
||||
|
||||
@pytest.fixture
|
||||
def registry(self):
|
||||
return FetchToolRegistry()
|
||||
|
||||
@pytest.fixture
|
||||
def mock_prisma(self):
|
||||
prisma = MagicMock()
|
||||
prisma.db = MagicMock()
|
||||
prisma.db.litellm_fetchtoolstable = MagicMock()
|
||||
return prisma
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_found(self, registry, mock_prisma):
|
||||
"""Test found tool."""
|
||||
from datetime import datetime, timezone
|
||||
now = datetime.now(timezone.utc)
|
||||
|
||||
mock_tool = MagicMock()
|
||||
mock_tool.__iter__ = MagicMock(
|
||||
return_value=iter(
|
||||
{
|
||||
"fetch_tool_id": "123",
|
||||
"fetch_tool_name": "test",
|
||||
"created_at": now,
|
||||
"updated_at": now,
|
||||
}.items()
|
||||
)
|
||||
)
|
||||
|
||||
mock_prisma.db.litellm_fetchtoolstable.find_unique = AsyncMock(return_value=mock_tool)
|
||||
|
||||
result = await registry.get_fetch_tool_from_db("123", mock_prisma)
|
||||
|
||||
assert result["fetch_tool_id"] == "123"
|
||||
assert result["fetch_tool_name"] == "test"
|
||||
assert "created_at" in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_not_found(self, registry, mock_prisma):
|
||||
"""Test not found."""
|
||||
mock_prisma.db.litellm_fetchtoolstable.find_unique = AsyncMock(return_value=None)
|
||||
|
||||
result = await registry.get_fetch_tool_from_db("missing", mock_prisma)
|
||||
|
||||
assert result is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_error(self, registry, mock_prisma):
|
||||
"""Test error handling."""
|
||||
mock_prisma.db.litellm_fetchtoolstable.find_unique = AsyncMock(
|
||||
side_effect=Exception("db error")
|
||||
)
|
||||
|
||||
with pytest.raises(Exception, match="db error"):
|
||||
await registry.get_fetch_tool_from_db("bad", mock_prisma)
|
||||
|
||||
|
||||
class TestGetAllFetchTools:
|
||||
"""Test get_all_fetch_tools_from_db."""
|
||||
|
||||
@pytest.fixture
|
||||
def registry(self):
|
||||
return FetchToolRegistry()
|
||||
|
||||
@pytest.fixture
|
||||
def mock_prisma(self):
|
||||
prisma = MagicMock()
|
||||
prisma.db = MagicMock()
|
||||
prisma.db.litellm_fetchtoolstable = MagicMock()
|
||||
return prisma
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_success(self, registry, mock_prisma):
|
||||
"""Test listing all tools."""
|
||||
from datetime import datetime, timezone
|
||||
now = datetime.now(timezone.utc)
|
||||
|
||||
mock_tool1 = MagicMock()
|
||||
mock_tool1.__iter__ = MagicMock(
|
||||
return_value=iter(
|
||||
{"fetch_tool_id": "1", "fetch_tool_name": "tool1", "created_at": now}.items()
|
||||
)
|
||||
)
|
||||
mock_tool2 = MagicMock()
|
||||
mock_tool2.__iter__ = MagicMock(
|
||||
return_value=iter(
|
||||
{"fetch_tool_id": "2", "fetch_tool_name": "tool2", "created_at": now}.items()
|
||||
)
|
||||
)
|
||||
|
||||
mock_prisma.db.litellm_fetchtoolstable.find_many = AsyncMock(
|
||||
return_value=[mock_tool1, mock_tool2]
|
||||
)
|
||||
|
||||
result = await registry.get_all_fetch_tools_from_db(mock_prisma)
|
||||
|
||||
assert len(result) == 2
|
||||
assert result[0]["fetch_tool_name"] == "tool1"
|
||||
assert result[1]["fetch_tool_name"] == "tool2"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty(self, registry, mock_prisma):
|
||||
"""Test empty list."""
|
||||
mock_prisma.db.litellm_fetchtoolstable.find_many = AsyncMock(return_value=[])
|
||||
|
||||
result = await registry.get_all_fetch_tools_from_db(mock_prisma)
|
||||
|
||||
assert result == []
|
||||
|
||||
|
||||
class TestUpdateFetchTool:
|
||||
"""Test update_fetch_tool_in_db."""
|
||||
|
||||
@pytest.fixture
|
||||
def registry(self):
|
||||
return FetchToolRegistry()
|
||||
|
||||
@pytest.fixture
|
||||
def mock_prisma(self):
|
||||
prisma = MagicMock()
|
||||
prisma.db = MagicMock()
|
||||
prisma.db.litellm_fetchtoolstable = MagicMock()
|
||||
return prisma
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_success(self, registry, mock_prisma):
|
||||
"""Test successful update."""
|
||||
mock_tool = MagicMock()
|
||||
mock_tool.fetch_tool_id = "123"
|
||||
mock_tool.created_at = MagicMock()
|
||||
mock_tool.created_at.isoformat.return_value = "2026-05-03T12:00:00"
|
||||
mock_tool.updated_at = MagicMock()
|
||||
mock_tool.updated_at.isoformat.return_value = "2026-05-03T13:00:00"
|
||||
|
||||
mock_prisma.db.litellm_fetchtoolstable.update = AsyncMock(return_value=mock_tool)
|
||||
|
||||
fetch_tool = {
|
||||
"fetch_tool_name": "updated-tool",
|
||||
"litellm_params": {"provider": "new"},
|
||||
}
|
||||
|
||||
result = await registry.update_fetch_tool_in_db("123", fetch_tool, mock_prisma)
|
||||
|
||||
assert result["fetch_tool_id"] == "123"
|
||||
assert result["fetch_tool_name"] == "updated-tool"
|
||||
mock_prisma.db.litellm_fetchtoolstable.update.assert_awaited_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_error(self, registry, mock_prisma):
|
||||
"""Test update error."""
|
||||
mock_prisma.db.litellm_fetchtoolstable.update = AsyncMock(
|
||||
side_effect=Exception("update failed")
|
||||
)
|
||||
|
||||
with pytest.raises(Exception, match="update failed"):
|
||||
await registry.update_fetch_tool_in_db("123", {}, mock_prisma)
|
||||
218
tests/webfetch_tests/test_firecrawl_fetch_extended.py
Normal file
218
tests/webfetch_tests/test_firecrawl_fetch_extended.py
Normal file
|
|
@ -0,0 +1,218 @@
|
|||
"""Tests for Firecrawl Fetch Transformation.
|
||||
|
||||
Covers FirecrawlFetchConfig scrape/fetch logic and self-hosted support.
|
||||
All network calls are mocked with httpx + respx.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import MagicMock, AsyncMock, patch
|
||||
|
||||
from litellm.llms.firecrawl.fetch.transformation import FirecrawlFetchConfig
|
||||
from litellm.llms.base_llm.fetch.transformation import WebFetchResponse
|
||||
|
||||
|
||||
class TestFirecrawlFetchConfig:
|
||||
"""Test FirecrawlFetchConfig initialization and methods."""
|
||||
|
||||
def test_ui_friendly_name(self):
|
||||
"""Test ui_friendly_name returns correct value."""
|
||||
config = FirecrawlFetchConfig()
|
||||
assert config.ui_friendly_name() == "Firecrawl"
|
||||
|
||||
def test_default_api_base(self):
|
||||
"""Test default API base is Firecrawl Cloud."""
|
||||
config = FirecrawlFetchConfig()
|
||||
assert "api.firecrawl.dev" in config.api_base
|
||||
|
||||
def test_custom_api_base(self):
|
||||
"""Test custom api_base overrides default."""
|
||||
config = FirecrawlFetchConfig(api_base="http://localhost:3002")
|
||||
assert config.api_base == "http://localhost:3002"
|
||||
|
||||
def test_validate_environment_no_api_key(self):
|
||||
"""Test validation without API key raises ValueError."""
|
||||
config = FirecrawlFetchConfig()
|
||||
with pytest.raises(ValueError, match="api_key"):
|
||||
config.validate_environment(headers={})
|
||||
|
||||
def test_validate_environment_with_api_key(self):
|
||||
"""Test validation with API key succeeds."""
|
||||
config = FirecrawlFetchConfig()
|
||||
headers = {"Authorization": "Bearer fc-test"}
|
||||
result = config.validate_environment(headers=headers, api_key="fc-test")
|
||||
assert "Authorization" in result
|
||||
|
||||
def test_scrape_url_sync_success(self):
|
||||
"""Test synchronous scrape URL success."""
|
||||
config = FirecrawlFetchConfig(api_key="fc-test")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {
|
||||
"success": True,
|
||||
"data": {
|
||||
"markdown": "# Hello\nWorld",
|
||||
"metadata": {
|
||||
"title": "Hello Page",
|
||||
"sourceURL": "https://example.com",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.post.return_value = mock_response
|
||||
|
||||
result = config.scrape_url(
|
||||
url="https://example.com",
|
||||
headers={},
|
||||
client=mock_client,
|
||||
)
|
||||
|
||||
assert isinstance(result, WebFetchResponse)
|
||||
assert result.url == "https://example.com"
|
||||
assert "Hello" in result.title
|
||||
|
||||
def test_scrape_url_sync_error(self):
|
||||
"""Test synchronous scrape URL error handling."""
|
||||
config = FirecrawlFetchConfig(api_key="fc-test")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 500
|
||||
mock_response.text = "Internal Server Error"
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.post.return_value = mock_response
|
||||
|
||||
with pytest.raises(ValueError, match="scrape URL"):
|
||||
config.scrape_url(
|
||||
url="https://example.com",
|
||||
headers={},
|
||||
client=mock_client,
|
||||
)
|
||||
|
||||
def test_scrape_url_sync_no_markdown(self):
|
||||
"""Test scrape URL with no markdown in response uses html."""
|
||||
config = FirecrawlFetchConfig(api_key="fc-test")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {
|
||||
"success": True,
|
||||
"data": {
|
||||
"html": "<html><body>Hello</body></html>",
|
||||
"metadata": {
|
||||
"title": "Hello",
|
||||
"sourceURL": "https://example.com",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.post.return_value = mock_response
|
||||
|
||||
result = config.scrape_url(
|
||||
url="https://example.com",
|
||||
headers={},
|
||||
client=mock_client,
|
||||
)
|
||||
|
||||
assert "Hello" in result.content
|
||||
|
||||
|
||||
class TestFetchUrlAsync:
|
||||
"""Test async fetch_url method."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_afetch_url_success(self):
|
||||
"""Test async fetch URL success."""
|
||||
config = FirecrawlFetchConfig(api_key="fc-test")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json = AsyncMock(return_value={
|
||||
"success": True,
|
||||
"data": {
|
||||
"markdown": "Async Hello",
|
||||
"metadata": {
|
||||
"title": "Async Hello",
|
||||
"sourceURL": "https://example.com",
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
result = await config.afetch_url(
|
||||
url="https://example.com",
|
||||
headers={},
|
||||
client=mock_client,
|
||||
)
|
||||
|
||||
assert isinstance(result, WebFetchResponse)
|
||||
assert result.content == "Async Hello"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_afetch_url_error(self):
|
||||
"""Test async fetch URL error."""
|
||||
config = FirecrawlFetchConfig(api_key="fc-test")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 429
|
||||
mock_response.text = "Rate limited"
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
with pytest.raises(ValueError, match="scrape URL"):
|
||||
await config.afetch_url(
|
||||
url="https://example.com",
|
||||
headers={},
|
||||
client=mock_client,
|
||||
)
|
||||
|
||||
|
||||
class TestSelfHosted:
|
||||
"""Test self-hosted Firecrawl instance support."""
|
||||
|
||||
def test_self_hosted_api_base(self):
|
||||
"""Test setting custom api_base for self-hosted instance."""
|
||||
config = FirecrawlFetchConfig(
|
||||
api_base="http://localhost:3002",
|
||||
api_key="local-key",
|
||||
)
|
||||
assert "localhost:3002" in config.api_base
|
||||
|
||||
def test_self_hosted_request_url(self):
|
||||
"""Test request URL uses custom api_base."""
|
||||
config = FirecrawlFetchConfig(
|
||||
api_base="http://my-firecrawl.internal:3002",
|
||||
api_key="local-key",
|
||||
)
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {
|
||||
"success": True,
|
||||
"data": {
|
||||
"markdown": "Local content",
|
||||
"metadata": {
|
||||
"title": "Local",
|
||||
"sourceURL": "https://example.com",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.post.return_value = mock_response
|
||||
|
||||
result = config.scrape_url(
|
||||
url="https://example.com",
|
||||
headers={},
|
||||
client=mock_client,
|
||||
)
|
||||
|
||||
# Verify the URL used in the request
|
||||
call_args = mock_client.post.call_args
|
||||
assert "my-firecrawl.internal:3002" in call_args[0][0] or \
|
||||
"my-firecrawl.internal:3002" in call_args[1].get("url", "")
|
||||
Loading…
Add table
Reference in a new issue