mirror of
https://github.com/mendableai/firecrawl.git
synced 2025-12-31 17:24:25 +00:00
278 lines
10 KiB
Python
278 lines
10 KiB
Python
"""
|
|
Integration tests for agent method with mocked requests.
|
|
"""
|
|
|
|
import unittest
|
|
from unittest.mock import patch, MagicMock
|
|
from pydantic import BaseModel, Field
|
|
from typing import List, Optional
|
|
|
|
from firecrawl import FirecrawlApp
|
|
|
|
|
|
class Founder(BaseModel):
|
|
name: str = Field(description="Full name of the founder")
|
|
role: Optional[str] = Field(None, description="Role or position")
|
|
background: Optional[str] = Field(None, description="Professional background")
|
|
|
|
|
|
class FoundersSchema(BaseModel):
|
|
founders: List[Founder] = Field(description="List of founders")
|
|
|
|
|
|
class TestAgent(unittest.TestCase):
|
|
"""Integration tests for agent method."""
|
|
|
|
@patch('firecrawl.v2.utils.http_client.requests.post')
|
|
@patch('firecrawl.v2.utils.http_client.requests.get')
|
|
def test_agent_basic(self, mock_get, mock_post):
|
|
"""Test basic agent call."""
|
|
# Mock start agent response
|
|
mock_start_response = MagicMock()
|
|
mock_start_response.ok = True
|
|
mock_start_response.status_code = 200
|
|
mock_start_response.json.return_value = {
|
|
"success": True,
|
|
"id": "test-agent-123",
|
|
"status": "processing"
|
|
}
|
|
mock_post.return_value = mock_start_response
|
|
|
|
# Mock get status response (completed)
|
|
mock_status_response = MagicMock()
|
|
mock_status_response.ok = True
|
|
mock_status_response.status_code = 200
|
|
mock_status_response.json.return_value = {
|
|
"success": True,
|
|
"id": "test-agent-123",
|
|
"status": "completed",
|
|
"data": {
|
|
"founders": [
|
|
{"name": "John Doe", "role": "CEO", "background": "Tech entrepreneur"},
|
|
{"name": "Jane Smith", "role": "CTO", "background": "Software engineer"}
|
|
]
|
|
},
|
|
"creditsUsed": 10,
|
|
"expiresAt": "2024-01-01T00:00:00Z"
|
|
}
|
|
mock_get.return_value = mock_status_response
|
|
|
|
app = FirecrawlApp(api_key="test-api-key")
|
|
result = app.agent(
|
|
prompt="Find the founders of Firecrawl",
|
|
schema=FoundersSchema
|
|
)
|
|
|
|
# Verify post was called with correct URL and data
|
|
mock_post.assert_called_once()
|
|
post_call_args = mock_post.call_args
|
|
post_url = post_call_args[1]["url"] if "url" in post_call_args[1] else post_call_args[0][0]
|
|
assert "/v2/agent" in str(post_url)
|
|
|
|
# Check request body
|
|
request_body = post_call_args[1]["json"]
|
|
assert request_body["prompt"] == "Find the founders of Firecrawl"
|
|
assert "schema" in request_body
|
|
assert request_body["schema"]["type"] == "object"
|
|
assert "founders" in request_body["schema"]["properties"]
|
|
|
|
# Verify get was called to check status
|
|
mock_get.assert_called()
|
|
|
|
# Check result
|
|
assert result.status == "completed"
|
|
assert result.data is not None
|
|
|
|
@patch('firecrawl.v2.utils.http_client.requests.post')
|
|
def test_agent_with_urls(self, mock_post):
|
|
"""Test agent call with URLs."""
|
|
mock_response = MagicMock()
|
|
mock_response.ok = True
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = {
|
|
"success": True,
|
|
"status": "completed",
|
|
"data": {"result": "done"}
|
|
}
|
|
mock_post.return_value = mock_response
|
|
|
|
app = FirecrawlApp(api_key="test-api-key")
|
|
result = app.agent(
|
|
urls=["https://example.com", "https://test.com"],
|
|
prompt="Extract information",
|
|
schema={"type": "object", "properties": {"info": {"type": "string"}}}
|
|
)
|
|
|
|
# Check request body includes URLs
|
|
post_call_args = mock_post.call_args
|
|
request_body = post_call_args[1]["json"]
|
|
assert request_body["urls"] == ["https://example.com", "https://test.com"]
|
|
assert request_body["prompt"] == "Extract information"
|
|
|
|
@patch('firecrawl.v2.utils.http_client.requests.post')
|
|
def test_agent_with_dict_schema(self, mock_post):
|
|
"""Test agent call with dict schema."""
|
|
mock_response = MagicMock()
|
|
mock_response.ok = True
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = {
|
|
"success": True,
|
|
"status": "completed",
|
|
"data": {"result": "done"}
|
|
}
|
|
mock_post.return_value = mock_response
|
|
|
|
schema = {
|
|
"type": "object",
|
|
"properties": {
|
|
"name": {"type": "string"},
|
|
"age": {"type": "integer"}
|
|
}
|
|
}
|
|
|
|
app = FirecrawlApp(api_key="test-api-key")
|
|
result = app.agent(
|
|
prompt="Extract person data",
|
|
schema=schema
|
|
)
|
|
|
|
# Check request body includes schema
|
|
post_call_args = mock_post.call_args
|
|
request_body = post_call_args[1]["json"]
|
|
assert request_body["schema"] == schema
|
|
|
|
@patch('firecrawl.v2.utils.http_client.requests.post')
|
|
def test_agent_with_all_params(self, mock_post):
|
|
"""Test agent call with all parameters."""
|
|
mock_response = MagicMock()
|
|
mock_response.ok = True
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = {
|
|
"success": True,
|
|
"status": "completed",
|
|
"data": {"result": "done"}
|
|
}
|
|
mock_post.return_value = mock_response
|
|
|
|
schema = {"type": "object"}
|
|
urls = ["https://example.com"]
|
|
|
|
app = FirecrawlApp(api_key="test-api-key")
|
|
result = app.agent(
|
|
urls=urls,
|
|
prompt="Complete test",
|
|
schema=schema,
|
|
integration="test-integration",
|
|
max_credits=50,
|
|
strict_constrain_to_urls=True,
|
|
poll_interval=1,
|
|
timeout=30
|
|
)
|
|
|
|
# Check all parameters are in request body
|
|
post_call_args = mock_post.call_args
|
|
request_body = post_call_args[1]["json"]
|
|
assert request_body["prompt"] == "Complete test"
|
|
assert request_body["urls"] == urls
|
|
assert request_body["schema"] == schema
|
|
assert request_body["integration"] == "test-integration"
|
|
assert request_body["maxCredits"] == 50
|
|
assert request_body["strictConstrainToURLs"] is True
|
|
|
|
@patch('firecrawl.v2.utils.http_client.requests.post')
|
|
def test_agent_pydantic_schema_normalization(self, mock_post):
|
|
"""Test that Pydantic schemas are properly normalized."""
|
|
mock_response = MagicMock()
|
|
mock_response.ok = True
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = {
|
|
"success": True,
|
|
"status": "completed",
|
|
"data": {"result": "done"}
|
|
}
|
|
mock_post.return_value = mock_response
|
|
|
|
app = FirecrawlApp(api_key="test-api-key")
|
|
result = app.agent(
|
|
prompt="Find founders",
|
|
schema=FoundersSchema
|
|
)
|
|
|
|
# Check that schema was normalized to JSON schema format
|
|
post_call_args = mock_post.call_args
|
|
request_body = post_call_args[1]["json"]
|
|
assert "schema" in request_body
|
|
schema = request_body["schema"]
|
|
assert schema["type"] == "object"
|
|
assert "properties" in schema
|
|
assert "founders" in schema["properties"]
|
|
assert schema["properties"]["founders"]["type"] == "array"
|
|
|
|
@patch('firecrawl.v2.utils.http_client.requests.post')
|
|
@patch('firecrawl.v2.utils.http_client.requests.get')
|
|
def test_agent_url_construction(self, mock_get, mock_post):
|
|
"""Test that agent requests are sent to correct URL."""
|
|
# Mock start agent response
|
|
mock_start_response = MagicMock()
|
|
mock_start_response.ok = True
|
|
mock_start_response.status_code = 200
|
|
mock_start_response.json.return_value = {
|
|
"success": True,
|
|
"id": "test-agent-123",
|
|
"status": "processing"
|
|
}
|
|
mock_post.return_value = mock_start_response
|
|
|
|
# Mock get status response
|
|
mock_status_response = MagicMock()
|
|
mock_status_response.ok = True
|
|
mock_status_response.status_code = 200
|
|
mock_status_response.json.return_value = {
|
|
"success": True,
|
|
"id": "test-agent-123",
|
|
"status": "completed",
|
|
"data": {"result": "done"}
|
|
}
|
|
mock_get.return_value = mock_status_response
|
|
|
|
app = FirecrawlApp(api_key="test-api-key", api_url="https://api.firecrawl.dev")
|
|
result = app.agent(prompt="Test prompt")
|
|
|
|
# Check POST URL - requests.post is called with url as keyword arg
|
|
post_call_args = mock_post.call_args
|
|
post_url = post_call_args[1].get("url") if "url" in post_call_args[1] else post_call_args[0][0]
|
|
assert "/v2/agent" in str(post_url)
|
|
|
|
# Check GET URL
|
|
get_call_args = mock_get.call_args
|
|
get_url = get_call_args[1].get("url") if "url" in get_call_args[1] else get_call_args[0][0]
|
|
assert "/v2/agent/test-agent-123" in str(get_url)
|
|
|
|
@patch('firecrawl.v2.utils.http_client.requests.post')
|
|
def test_agent_headers(self, mock_post):
|
|
"""Test that agent requests include correct headers."""
|
|
mock_response = MagicMock()
|
|
mock_response.ok = True
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = {
|
|
"success": True,
|
|
"status": "completed",
|
|
"data": {"result": "done"}
|
|
}
|
|
mock_post.return_value = mock_response
|
|
|
|
app = FirecrawlApp(api_key="test-api-key")
|
|
result = app.agent(prompt="Test prompt")
|
|
|
|
# Check headers
|
|
post_call_args = mock_post.call_args
|
|
headers = post_call_args[1]["headers"]
|
|
assert "Authorization" in headers
|
|
assert headers["Authorization"] == "Bearer test-api-key"
|
|
assert headers["Content-Type"] == "application/json"
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|
|
|