firecrawl 4.12.0__py3-none-any.whl

firecrawl/__tests__/unit/v2/methods/test_agent.py

@@ -0,0 +1,367 @@
+"""
+Unit tests for agent methods with mocked HTTP client.
+"""
+
+import pytest
+import time
+from unittest.mock import Mock, patch
+from pydantic import BaseModel, Field
+from typing import List, Optional
+
+from firecrawl.v2.methods.agent import (
+    start_agent,
+    agent,
+    get_agent_status,
+    cancel_agent,
+    wait_agent
+)
+from firecrawl.v2.types import AgentResponse
+from firecrawl.v2.utils.error_handler import BadRequestError
+
+
+class TestAgentMethods:
+    """Unit tests for agent methods with mocked HTTP client."""
+
+    def setup_method(self):
+        """Set up test fixtures."""
+        self.mock_client = Mock()
+        self.job_id = "test-agent-123"
+        
+        # Sample agent response
+        self.sample_response = {
+            "success": True,
+            "id": self.job_id,
+            "status": "completed",
+            "data": {
+                "founders": [
+                    {"name": "John Doe", "role": "CEO"},
+                    {"name": "Jane Smith", "role": "CTO"}
+                ]
+            },
+            "creditsUsed": 10,
+            "expiresAt": "2024-01-01T00:00:00Z"
+        }
+
+    def test_start_agent_basic(self):
+        """Test starting an agent job with basic parameters."""
+        mock_response = Mock()
+        mock_response.ok = True
+        mock_response.json.return_value = {
+            "success": True,
+            "id": self.job_id,
+            "status": "processing"
+        }
+        
+        self.mock_client.post.return_value = mock_response
+        
+        result = start_agent(
+            self.mock_client,
+            None,
+            prompt="Find information about Firecrawl"
+        )
+        
+        # Check that post was called with correct endpoint
+        self.mock_client.post.assert_called_once()
+        call_args = self.mock_client.post.call_args
+        assert call_args[0][0] == "/v2/agent"
+        
+        # Check request body (second positional argument)
+        body = call_args[0][1]
+        assert body["prompt"] == "Find information about Firecrawl"
+        assert "urls" not in body
+        
+        # Check result
+        assert isinstance(result, AgentResponse)
+        assert result.id == self.job_id
+        assert result.status == "processing"
+
+    def test_start_agent_with_urls(self):
+        """Test starting an agent job with URLs."""
+        mock_response = Mock()
+        mock_response.ok = True
+        mock_response.json.return_value = {
+            "success": True,
+            "id": self.job_id,
+            "status": "processing"
+        }
+        
+        self.mock_client.post.return_value = mock_response
+        
+        urls = ["https://example.com", "https://test.com"]
+        result = start_agent(
+            self.mock_client,
+            urls,
+            prompt="Extract data"
+        )
+        
+        call_args = self.mock_client.post.call_args
+        body = call_args[0][1]
+        assert body["urls"] == urls
+
+    def test_start_agent_with_dict_schema(self):
+        """Test starting an agent job with dict schema."""
+        mock_response = Mock()
+        mock_response.ok = True
+        mock_response.json.return_value = {
+            "success": True,
+            "id": self.job_id,
+            "status": "processing"
+        }
+        
+        self.mock_client.post.return_value = mock_response
+        
+        schema = {
+            "type": "object",
+            "properties": {
+                "name": {"type": "string"}
+            }
+        }
+        
+        result = start_agent(
+            self.mock_client,
+            None,
+            prompt="Extract data",
+            schema=schema
+        )
+        
+        call_args = self.mock_client.post.call_args
+        body = call_args[0][1]
+        assert body["schema"] == schema
+
+    def test_start_agent_with_pydantic_schema(self):
+        """Test starting an agent job with Pydantic schema."""
+        mock_response = Mock()
+        mock_response.ok = True
+        mock_response.json.return_value = {
+            "success": True,
+            "id": self.job_id,
+            "status": "processing"
+        }
+        
+        self.mock_client.post.return_value = mock_response
+        
+        class Founder(BaseModel):
+            name: str = Field(description="Full name")
+            role: Optional[str] = Field(None, description="Role")
+        
+        class FoundersSchema(BaseModel):
+            founders: List[Founder] = Field(description="List of founders")
+        
+        result = start_agent(
+            self.mock_client,
+            None,
+            prompt="Find founders",
+            schema=FoundersSchema
+        )
+        
+        call_args = self.mock_client.post.call_args
+        body = call_args[0][1]
+        assert "schema" in body
+        assert body["schema"]["type"] == "object"
+        assert "founders" in body["schema"]["properties"]
+
+    def test_start_agent_with_all_params(self):
+        """Test starting an agent job with all parameters."""
+        mock_response = Mock()
+        mock_response.ok = True
+        mock_response.json.return_value = {
+            "success": True,
+            "id": self.job_id,
+            "status": "processing"
+        }
+        
+        self.mock_client.post.return_value = mock_response
+        
+        schema = {"type": "object"}
+        urls = ["https://example.com"]
+        
+        result = start_agent(
+            self.mock_client,
+            urls,
+            prompt="Complete test",
+            schema=schema,
+            integration="test-integration",
+            max_credits=50,
+            strict_constrain_to_urls=True
+        )
+        
+        call_args = self.mock_client.post.call_args
+        body = call_args[0][1]
+        assert body["prompt"] == "Complete test"
+        assert body["urls"] == urls
+        assert body["schema"] == schema
+        assert body["integration"] == "test-integration"
+        assert body["maxCredits"] == 50
+        assert body["strictConstrainToURLs"] is True
+
+    def test_get_agent_status(self):
+        """Test getting agent status."""
+        mock_response = Mock()
+        mock_response.ok = True
+        mock_response.json.return_value = self.sample_response
+        
+        self.mock_client.get.return_value = mock_response
+        
+        result = get_agent_status(self.mock_client, self.job_id)
+        
+        # Check that get was called with correct endpoint
+        self.mock_client.get.assert_called_once_with(f"/v2/agent/{self.job_id}")
+        
+        # Check result
+        assert isinstance(result, AgentResponse)
+        assert result.id == self.job_id
+        assert result.status == "completed"
+
+    def test_cancel_agent(self):
+        """Test canceling an agent job."""
+        mock_response = Mock()
+        mock_response.ok = True
+        mock_response.json.return_value = {"success": True}
+        
+        self.mock_client.delete.return_value = mock_response
+        
+        result = cancel_agent(self.mock_client, self.job_id)
+        
+        # Check that delete was called with correct endpoint
+        self.mock_client.delete.assert_called_once_with(f"/v2/agent/{self.job_id}")
+        
+        # Check result
+        assert result is True
+
+    @patch('time.sleep')
+    def test_wait_agent_completed(self, mock_sleep):
+        """Test waiting for agent to complete."""
+        # First call returns processing, second returns completed
+        mock_response_processing = Mock()
+        mock_response_processing.ok = True
+        mock_response_processing.json.return_value = {
+            "success": True,
+            "id": self.job_id,
+            "status": "processing"
+        }
+        
+        mock_response_completed = Mock()
+        mock_response_completed.ok = True
+        mock_response_completed.json.return_value = self.sample_response
+        
+        self.mock_client.get.side_effect = [
+            mock_response_processing,
+            mock_response_completed
+        ]
+        
+        result = wait_agent(self.mock_client, self.job_id, poll_interval=1)
+        
+        # Should have called get twice
+        assert self.mock_client.get.call_count == 2
+        assert result.status == "completed"
+        assert mock_sleep.call_count == 1
+
+    @patch('time.sleep')
+    def test_wait_agent_timeout(self, mock_sleep):
+        """Test waiting for agent with timeout."""
+        mock_response = Mock()
+        mock_response.ok = True
+        mock_response.json.return_value = {
+            "success": True,
+            "id": self.job_id,
+            "status": "processing"
+        }
+        
+        self.mock_client.get.return_value = mock_response
+        
+        # Mock time.time to simulate timeout
+        with patch('time.time', side_effect=[0, 0, 5]):  # Start at 0, timeout at 5
+            result = wait_agent(
+                self.mock_client,
+                self.job_id,
+                poll_interval=1,
+                timeout=3  # Timeout after 3 seconds
+            )
+        
+        # Should return processing status due to timeout
+        assert result.status == "processing"
+
+    @patch('time.sleep')
+    def test_agent_complete_flow(self, mock_sleep):
+        """Test the complete agent flow (start + wait)."""
+        # Mock start_agent response
+        mock_start_response = Mock()
+        mock_start_response.ok = True
+        mock_start_response.json.return_value = {
+            "success": True,
+            "id": self.job_id,
+            "status": "processing"
+        }
+        
+        # Mock get_agent_status responses
+        mock_status_response = Mock()
+        mock_status_response.ok = True
+        mock_status_response.json.return_value = self.sample_response
+        
+        self.mock_client.post.return_value = mock_start_response
+        self.mock_client.get.return_value = mock_status_response
+        
+        result = agent(
+            self.mock_client,
+            None,
+            prompt="Find information",
+            poll_interval=1
+        )
+        
+        # Should have called post once and get once
+        assert self.mock_client.post.call_count == 1
+        assert self.mock_client.get.call_count == 1
+        
+        # Check result
+        assert isinstance(result, AgentResponse)
+        assert result.status == "completed"
+        assert result.data is not None
+
+    def test_agent_immediate_completion(self):
+        """Test agent that completes immediately (no job ID)."""
+        mock_response = Mock()
+        mock_response.ok = True
+        mock_response.json.return_value = {
+            "success": True,
+            "status": "completed",
+            "data": {"result": "done"}
+        }
+        
+        self.mock_client.post.return_value = mock_response
+        
+        result = agent(
+            self.mock_client,
+            None,
+            prompt="Quick task"
+        )
+        
+        # Should only call post, not get
+        assert self.mock_client.post.call_count == 1
+        assert self.mock_client.get.call_count == 0
+        assert result.status == "completed"
+
+    def test_start_agent_error_handling(self):
+        """Test error handling in start_agent."""
+        mock_response = Mock()
+        mock_response.ok = False
+        mock_response.status_code = 400
+        mock_response.text = "Bad Request"
+        # Mock response.json() to return error details
+        mock_response.json.return_value = {
+            "error": "Invalid request",
+            "details": "Bad Request"
+        }
+        
+        self.mock_client.post.return_value = mock_response
+        
+        with pytest.raises(BadRequestError) as exc_info:
+            start_agent(
+                self.mock_client,
+                None,
+                prompt="Test prompt"
+            )
+        
+        # Verify the exception has the correct status code
+        assert exc_info.value.status_code == 400
+        assert "agent" in str(exc_info.value).lower()
+

firecrawl/__tests__/unit/v2/methods/test_agent_request_preparation.py

@@ -0,0 +1,226 @@
+"""
+Unit tests for agent request preparation.
+"""
+
+import pytest
+from pydantic import BaseModel, Field
+from typing import List, Optional
+
+from firecrawl.v2.methods.agent import _prepare_agent_request
+
+
+class TestAgentRequestPreparation:
+    """Unit tests for agent request preparation."""
+
+    def test_basic_request_preparation(self):
+        """Test basic request preparation with minimal fields."""
+        data = _prepare_agent_request(
+            None,
+            prompt="Find information about Firecrawl"
+        )
+        
+        assert data["prompt"] == "Find information about Firecrawl"
+        assert "urls" not in data
+        assert "schema" not in data
+
+    def test_request_with_urls(self):
+        """Test request preparation with URLs."""
+        urls = ["https://example.com", "https://test.com"]
+        data = _prepare_agent_request(
+            urls,
+            prompt="Extract data from these pages"
+        )
+        
+        assert data["prompt"] == "Extract data from these pages"
+        assert data["urls"] == urls
+
+    def test_request_with_dict_schema(self):
+        """Test request preparation with dict schema."""
+        schema = {
+            "type": "object",
+            "properties": {
+                "name": {"type": "string"},
+                "age": {"type": "integer"}
+            }
+        }
+        data = _prepare_agent_request(
+            None,
+            prompt="Extract person data",
+            schema=schema
+        )
+        
+        assert data["prompt"] == "Extract person data"
+        assert data["schema"] == schema
+
+    def test_request_with_pydantic_schema(self):
+        """Test request preparation with Pydantic BaseModel schema."""
+        class Person(BaseModel):
+            name: str = Field(description="Person's name")
+            age: Optional[int] = Field(None, description="Person's age")
+        
+        data = _prepare_agent_request(
+            None,
+            prompt="Extract person data",
+            schema=Person
+        )
+        
+        assert data["prompt"] == "Extract person data"
+        assert "schema" in data
+        assert data["schema"]["type"] == "object"
+        assert "properties" in data["schema"]
+        assert "name" in data["schema"]["properties"]
+        assert "age" in data["schema"]["properties"]
+
+    def test_request_with_pydantic_schema_instance(self):
+        """Test request preparation with Pydantic model instance."""
+        class Person(BaseModel):
+            name: str = Field(description="Person's name")
+            age: Optional[int] = Field(None, description="Person's age")
+        
+        person_instance = Person(name="John", age=30)
+        data = _prepare_agent_request(
+            None,
+            prompt="Extract person data",
+            schema=person_instance
+        )
+        
+        assert data["prompt"] == "Extract person data"
+        assert "schema" in data
+        # Should use the class schema, not the instance data
+        assert data["schema"]["type"] == "object"
+
+    def test_request_with_nested_pydantic_schema(self):
+        """Test request preparation with nested Pydantic schema."""
+        class Founder(BaseModel):
+            name: str = Field(description="Full name of the founder")
+            role: Optional[str] = Field(None, description="Role or position")
+        
+        class FoundersSchema(BaseModel):
+            founders: List[Founder] = Field(description="List of founders")
+        
+        data = _prepare_agent_request(
+            None,
+            prompt="Find the founders",
+            schema=FoundersSchema
+        )
+        
+        assert data["prompt"] == "Find the founders"
+        assert "schema" in data
+        assert data["schema"]["type"] == "object"
+        assert "founders" in data["schema"]["properties"]
+        assert data["schema"]["properties"]["founders"]["type"] == "array"
+
+    def test_request_with_integration(self):
+        """Test request preparation with integration tag."""
+        data = _prepare_agent_request(
+            None,
+            prompt="Test prompt",
+            integration="  test-integration  "
+        )
+        
+        assert data["prompt"] == "Test prompt"
+        assert data["integration"] == "test-integration"
+
+    def test_request_with_max_credits(self):
+        """Test request preparation with max credits."""
+        data = _prepare_agent_request(
+            None,
+            prompt="Test prompt",
+            max_credits=100
+        )
+        
+        assert data["prompt"] == "Test prompt"
+        assert data["maxCredits"] == 100
+
+    def test_request_with_strict_constrain_to_urls(self):
+        """Test request preparation with strict_constrain_to_urls."""
+        data = _prepare_agent_request(
+            ["https://example.com"],
+            prompt="Test prompt",
+            strict_constrain_to_urls=True
+        )
+        
+        assert data["prompt"] == "Test prompt"
+        assert data["strictConstrainToURLs"] is True
+
+    def test_request_all_fields(self):
+        """Test request preparation with all fields."""
+        schema = {
+            "type": "object",
+            "properties": {"test": {"type": "string"}}
+        }
+        urls = ["https://example.com"]
+        
+        data = _prepare_agent_request(
+            urls,
+            prompt="Complete test",
+            schema=schema,
+            integration="test-integration",
+            max_credits=50,
+            strict_constrain_to_urls=True
+        )
+        
+        assert data["prompt"] == "Complete test"
+        assert data["urls"] == urls
+        assert data["schema"] == schema
+        assert data["integration"] == "test-integration"
+        assert data["maxCredits"] == 50
+        assert data["strictConstrainToURLs"] is True
+
+    def test_request_with_empty_integration(self):
+        """Test that empty integration is not included."""
+        data = _prepare_agent_request(
+            None,
+            prompt="Test prompt",
+            integration="   "
+        )
+        
+        assert "integration" not in data
+
+    def test_request_with_zero_max_credits(self):
+        """Test that zero max_credits is not included."""
+        data = _prepare_agent_request(
+            None,
+            prompt="Test prompt",
+            max_credits=0
+        )
+        
+        assert "maxCredits" not in data
+
+    def test_request_with_false_strict_constrain(self):
+        """Test that False strict_constrain_to_urls is not included."""
+        data = _prepare_agent_request(
+            None,
+            prompt="Test prompt",
+            strict_constrain_to_urls=False
+        )
+        
+        assert "strictConstrainToURLs" not in data
+
+    def test_request_with_invalid_schema_type_string(self):
+        """Test that invalid schema types raise ValueError."""
+        with pytest.raises(ValueError, match="Invalid schema type"):
+            _prepare_agent_request(
+                None,
+                prompt="Test prompt",
+                schema="invalid_string_schema"
+            )
+
+    def test_request_with_invalid_schema_type_int(self):
+        """Test that invalid schema types raise ValueError."""
+        with pytest.raises(ValueError, match="Invalid schema type"):
+            _prepare_agent_request(
+                None,
+                prompt="Test prompt",
+                schema=123
+            )
+
+    def test_request_with_invalid_schema_type_list(self):
+        """Test that invalid schema types raise ValueError."""
+        with pytest.raises(ValueError, match="Invalid schema type"):
+            _prepare_agent_request(
+                None,
+                prompt="Test prompt",
+                schema=["not", "a", "valid", "schema"]
+            )
+

firecrawl/__tests__/unit/v2/methods/test_batch_request_preparation.py

@@ -0,0 +1,90 @@
+import pytest
+from firecrawl.v2.types import ScrapeOptions, Location, WebhookConfig
+from firecrawl.v2.methods.batch import prepare_batch_scrape_request
+
+
+class TestBatchScrapeRequestPreparation:
+    """Unit tests for batch scrape request preparation."""
+
+    def test_urls_validation(self):
+        # empty list
+        with pytest.raises(ValueError):
+            prepare_batch_scrape_request([])
+        # invalid protocol
+        with pytest.raises(ValueError):
+            prepare_batch_scrape_request(["example.com"])  # missing http(s)
+        # valid
+        data = prepare_batch_scrape_request(["https://example.com", "http://foo.bar"])
+        assert data["urls"] == ["https://example.com", "http://foo.bar"]
+
+    def test_flatten_scrape_options(self):
+        opts = ScrapeOptions(
+            formats=["markdown", "change_tracking", {"type": "screenshot", "full_page": True, "quality": 80}],
+            include_tags=["main"],
+            exclude_tags=["nav"],
+            only_main_content=True,
+            wait_for=500,
+            timeout=30000,
+            mobile=True,
+            parsers=["pdf"],
+            actions=[{"type": "screenshot", "full_page": True}],
+            location=Location(country="us", languages=["en"]),
+            skip_tls_verification=False,
+            remove_base64_images=False,
+            fast_mode=True,
+            use_mock="test",
+            block_ads=False,
+            proxy="basic",
+            max_age=1000,
+            store_in_cache=False,
+        )
+        data = prepare_batch_scrape_request(["https://example.com"], options=opts)
+
+        # Formats should be at top-level as list, with screenshot normalized to object w/ fullPage
+        assert isinstance(data.get("formats"), list)
+        assert "markdown" in data["formats"]
+        # snake_case format should be converted to camelCase
+        assert "changeTracking" in data["formats"]
+        found_obj = next((f for f in data["formats"] if isinstance(f, dict) and f.get("type") == "screenshot"), None)
+        assert found_obj is not None and found_obj.get("fullPage") is True and found_obj.get("quality") == 80
+
+        # Field conversions to camelCase
+        assert data["includeTags"] == ["main"]
+        assert data["excludeTags"] == ["nav"]
+        assert data["onlyMainContent"] is True
+        assert data["waitFor"] == 500
+        assert data["timeout"] == 30000
+        assert data["mobile"] is True
+        assert data["parsers"] == ["pdf"]
+        assert isinstance(data["actions"], list) and data["actions"][0]["type"] == "screenshot"
+        assert isinstance(data["location"], dict) and data["location"]["country"] == "us"
+        assert data["skipTlsVerification"] is False
+        assert data["removeBase64Images"] is False
+        assert data["fastMode"] is True
+        assert data["useMock"] == "test"
+        assert data["blockAds"] is False
+        assert data["proxy"] == "basic"
+        assert data["maxAge"] == 1000
+        assert data["storeInCache"] is False
+
+    def test_batch_specific_fields(self):
+        webhook = WebhookConfig(url="https://hook.test", headers={"X": "Y"}, events=["completed"]) 
+        data = prepare_batch_scrape_request(
+            ["https://example.com"],
+            webhook=webhook,
+            append_to_id="00000000-0000-0000-0000-000000000000",
+            ignore_invalid_urls=True,
+            max_concurrency=5,
+            zero_data_retention=True,
+            integration="_unit-test",
+        )
+        assert isinstance(data["webhook"], dict) and data["webhook"]["url"] == "https://hook.test"
+        assert data["appendToId"] == "00000000-0000-0000-0000-000000000000"
+        assert data["ignoreInvalidURLs"] is True
+        assert data["maxConcurrency"] == 5
+        assert data["zeroDataRetention"] is True
+        assert data["integration"] == "_unit-test"
+
+    def test_string_webhook_is_passed_verbatim(self):
+        data = prepare_batch_scrape_request(["https://example.com"], webhook="https://hook.simple")
+        assert data["webhook"] == "https://hook.simple"