fixing more tests

Author: filip-michalsky

stagehand/handlers/extract_handler.py

@@ -7,7 +7,7 @@
 from stagehand.a11y.utils import get_accessibility_tree
 from stagehand.llm.inference import extract as extract_inference
 from stagehand.metrics import StagehandFunctionName  # Changed import location
-from stagehand.types import DefaultExtractSchema, ExtractOptions, ExtractResult
+from stagehand.schemas import DEFAULT_EXTRACT_SCHEMA as DefaultExtractSchema, ExtractOptions, ExtractResult
 from stagehand.utils import inject_urls, transform_url_strings_to_ids
 
 T = TypeVar("T", bound=BaseModel)
@@ -153,10 +153,12 @@ async def extract(
                     f"Failed to validate extracted data against schema {schema.__name__}: {e}. Keeping raw data dict in .data field."
                 )
 
-        # Create ExtractResult object
-        result = ExtractResult(
-            data=processed_data_payload,
-        )
+        # Create ExtractResult object with extracted data as fields
+        if isinstance(processed_data_payload, dict):
+            result = ExtractResult(**processed_data_payload)
+        else:
+            # For non-dict data (like Pydantic models), create with data field
+            result = ExtractResult(data=processed_data_payload)
 
         return result
 

tests/conftest.py

@@ -85,8 +85,65 @@ def mock_stagehand_page(mock_playwright_page):
     mock_client.logger.error = MagicMock()
     mock_client._get_lock_for_session = MagicMock(return_value=AsyncMock())
     mock_client._execute = AsyncMock()
+    mock_client.update_metrics = MagicMock()
 
     stagehand_page = StagehandPage(mock_playwright_page, mock_client)
+    
+    # Mock CDP calls for accessibility tree
+    async def mock_send_cdp(method, params=None):
+        if method == "Accessibility.getFullAXTree":
+            return {
+                "nodes": [
+                    {
+                        "nodeId": "1",
+                        "role": {"value": "button"},
+                        "name": {"value": "Click me"},
+                        "backendDOMNodeId": 1,
+                        "childIds": [],
+                        "properties": []
+                    },
+                    {
+                        "nodeId": "2", 
+                        "role": {"value": "textbox"},
+                        "name": {"value": "Search input"},
+                        "backendDOMNodeId": 2,
+                        "childIds": [],
+                        "properties": []
+                    }
+                ]
+            }
+        elif method == "DOM.resolveNode":
+            return {
+                "object": {
+                    "objectId": "test-object-id"
+                }
+            }
+        elif method == "Runtime.callFunctionOn":
+            return {
+                "result": {
+                    "value": "//div[@id='test']"
+                }
+            }
+        return {}
+    
+    stagehand_page.send_cdp = AsyncMock(side_effect=mock_send_cdp)
+    
+    # Mock get_cdp_client to return a mock CDP session
+    mock_cdp_client = AsyncMock()
+    mock_cdp_client.send = AsyncMock(return_value={"result": {"value": "//div[@id='test']"}})
+    stagehand_page.get_cdp_client = AsyncMock(return_value=mock_cdp_client)
+    
+    # Mock ensure_injection and evaluate methods
+    stagehand_page.ensure_injection = AsyncMock()
+    stagehand_page.evaluate = AsyncMock(return_value=[])
+    
+    # Mock enable/disable CDP domain methods
+    stagehand_page.enable_cdp_domain = AsyncMock()
+    stagehand_page.disable_cdp_domain = AsyncMock()
+    
+    # Mock _wait_for_settled_dom to avoid asyncio.sleep issues
+    stagehand_page._wait_for_settled_dom = AsyncMock()
+    
     return stagehand_page
 
 

tests/mocks/mock_llm.py

@@ -144,7 +144,9 @@ def _create_response(self, data: Any, model: str) -> MockLLMResponse:
         if isinstance(data, str):
             return MockLLMResponse(data, model=model)
         elif isinstance(data, dict):
-            content = data.get("content", str(data))
+            # For extract responses, convert dict to JSON string for content
+            import json
+            content = json.dumps(data)
             return MockLLMResponse(content, data=data, model=model)
         else:
             return MockLLMResponse(str(data), data=data, model=model)
@@ -247,4 +249,60 @@ def get_usage_stats(self) -> Dict[str, int]:
             "total_prompt_tokens": total_prompt_tokens,
             "total_completion_tokens": total_completion_tokens,
             "total_tokens": total_prompt_tokens + total_completion_tokens
-        } 
+        }
+    
+    def create_response(
+        self,
+        *,
+        messages: list[dict[str, str]],
+        model: Optional[str] = None,
+        function_name: Optional[str] = None,
+        **kwargs
+    ) -> MockLLMResponse:
+        """Create a response using the same interface as the real LLMClient"""
+        # Use function_name to determine response type if available
+        if function_name:
+            response_type = function_name.lower()
+        else:
+            # Fall back to content-based detection
+            content = str(messages).lower()
+            response_type = self._determine_response_type(content)
+        
+        # Track the call
+        self.call_count += 1
+        self.last_messages = messages
+        self.last_model = model or self.default_model
+        self.last_kwargs = kwargs
+        
+        # Store call in history
+        call_info = {
+            "messages": messages,
+            "model": self.last_model,
+            "kwargs": kwargs,
+            "function_name": function_name,
+            "timestamp": asyncio.get_event_loop().time()
+        }
+        self.call_history.append(call_info)
+        
+        # Simulate failure if configured
+        if self.should_fail:
+            raise Exception(self.failure_message)
+        
+        # Check for custom responses first
+        if response_type in self.custom_responses:
+            response_data = self.custom_responses[response_type]
+            if callable(response_data):
+                response_data = response_data(messages, **kwargs)
+            return self._create_response(response_data, model=self.last_model)
+        
+        # Use default response mapping
+        response_generator = self.response_mapping.get(response_type, self._default_response)
+        response_data = response_generator(messages, **kwargs)
+        
+        response = self._create_response(response_data, model=self.last_model)
+        
+        # Call metrics callback if set
+        if self.metrics_callback:
+            self.metrics_callback(response, 100, response_type)  # 100ms mock inference time
+        
+        return response