fixing tests

Author: filip-michalsky

tests/conftest.py

@@ -113,15 +113,44 @@ async def mock_send_cdp(method, params=None):
                 ]
             }
         elif method == "DOM.resolveNode":
+            # Create a mapping of element IDs to appropriate object IDs
+            backend_node_id = params.get("backendNodeId", 1)
             return {
                 "object": {
-                    "objectId": "test-object-id"
+                    "objectId": f"test-object-id-{backend_node_id}"
                 }
             }
         elif method == "Runtime.callFunctionOn":
+            # Map object IDs to appropriate selectors based on the element ID
+            object_id = params.get("objectId", "")
+            
+            # Extract backend_node_id from object_id
+            if "test-object-id-" in object_id:
+                backend_node_id = object_id.replace("test-object-id-", "")
+                
+                # Map specific element IDs to expected selectors for tests
+                selector_mapping = {
+                    "100": "//a[@id='home-link']",
+                    "101": "//a[@id='about-link']", 
+                    "102": "//a[@id='contact-link']",
+                    "200": "//button[@id='visible-button']",
+                    "300": "//input[@id='form-input']",
+                    "400": "//div[@id='target-element']",
+                    "501": "//button[@id='btn1']",
+                    "600": "//button[@id='interactive-btn']",
+                    "700": "//div[@id='positioned-element']",
+                    "800": "//div[@id='highlighted-element']",
+                    "900": "//div[@id='custom-model-element']",
+                    "1000": "//input[@id='complex-element']",
+                }
+                
+                xpath = selector_mapping.get(backend_node_id, "//div[@id='test']")
+            else:
+                xpath = "//div[@id='test']"
+                
             return {
                 "result": {
-                    "value": "//div[@id='test']"
+                    "value": xpath
                 }
             }
         return {}
@@ -130,7 +159,45 @@ async def mock_send_cdp(method, params=None):
 
     # 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']"}})
+    
+    # Set up the mock CDP client to handle Runtime.callFunctionOn properly
+    async def mock_cdp_send(method, params=None):
+        if method == "Runtime.callFunctionOn":
+            # Map object IDs to appropriate selectors based on the element ID
+            object_id = params.get("objectId", "")
+            
+            # Extract backend_node_id from object_id
+            if "test-object-id-" in object_id:
+                backend_node_id = object_id.replace("test-object-id-", "")
+                
+                # Map specific element IDs to expected selectors for tests
+                selector_mapping = {
+                    "100": "//a[@id='home-link']",
+                    "101": "//a[@id='about-link']", 
+                    "102": "//a[@id='contact-link']",
+                    "200": "//button[@id='visible-button']",
+                    "300": "//input[@id='form-input']",
+                    "400": "//div[@id='target-element']",
+                    "501": "//button[@id='btn1']",
+                    "600": "//button[@id='interactive-btn']",
+                    "700": "//div[@id='positioned-element']",
+                    "800": "//div[@id='highlighted-element']",
+                    "900": "//div[@id='custom-model-element']",
+                    "1000": "//input[@id='complex-element']",
+                }
+                
+                xpath = selector_mapping.get(backend_node_id, "//div[@id='test']")
+            else:
+                xpath = "//div[@id='test']"
+                
+            return {
+                "result": {
+                    "value": xpath
+                }
+            }
+        return {"result": {"value": "//div[@id='test']"}}
+    
+    mock_cdp_client.send = AsyncMock(side_effect=mock_cdp_send)
     stagehand_page.get_cdp_client = AsyncMock(return_value=mock_cdp_client)
 
     # Mock ensure_injection and evaluate methods

tests/integration/end_to_end/test_workflows.py

@@ -576,7 +576,7 @@ async def mock_extract(instruction, **kwargs):
                 # Extract data via Browserbase
                 extracted = await stagehand.page.extract("extract page title and content")
                 assert extracted["title"] == "Remote Page Title"
-                assert "Browserbase" in extracted["content"]
+                assert extracted["content"] == "Content extracted via Browserbase"
 
                 # Verify server interactions
                 assert server.was_called_with_endpoint("act")

tests/mocks/mock_llm.py

@@ -148,6 +148,13 @@ def _create_response(self, data: Any, model: str) -> MockLLMResponse:
             import json
             content = json.dumps(data)
             return MockLLMResponse(content, data=data, model=model)
+        elif isinstance(data, list):
+            # For observe responses, convert list to JSON string for content
+            import json
+            # Wrap the list in the expected format for observe responses
+            response_dict = {"elements": data}
+            content = json.dumps(response_dict)
+            return MockLLMResponse(content, data=response_dict, model=model)
         else:
             return MockLLMResponse(str(data), data=data, model=model)
 

tests/mocks/mock_server.py

@@ -199,21 +199,23 @@ def _extract_endpoint(self, url: str) -> str:
         # Remove base URL and extract the last path component
         path = url.split("/")[-1]
 
-        # Handle common Stagehand endpoints
+        # Handle common Stagehand endpoints - check exact matches to avoid substring issues
         if "session" in url and "create" in url:
-            return "create_session"
-        elif "navigate" in path:
-            return "navigate"
-        elif "act" in path:
-            return "act"
-        elif "observe" in path:
-            return "observe"
-        elif "extract" in path:
-            return "extract"
-        elif "screenshot" in path:
-            return "screenshot"
+            endpoint = "create_session"
+        elif path == "navigate":
+            endpoint = "navigate"
+        elif path == "act":
+            endpoint = "act"
+        elif path == "observe":
+            endpoint = "observe"
+        elif path == "extract":
+            endpoint = "extract"
+        elif path == "screenshot":
+            endpoint = "screenshot"
         else:
-            return path or "unknown"
+            endpoint = path or "unknown"
+            
+        return endpoint
 
     def set_response_override(self, endpoint: str, response: Union[Dict, callable]):
         """Override the default response for a specific endpoint"""