updates to integration tests

Author: filip-michalsky

stagehand/handlers/extract_handler.py

@@ -159,14 +159,9 @@ async def extract(
                 )
 
         # Create ExtractResult object with extracted data as fields
-        if isinstance(processed_data_payload, dict):
-            result = ExtractResult(**processed_data_payload)
-        elif hasattr(processed_data_payload, "model_dump"):
-            # For Pydantic models, convert to dict and spread as fields
-            result = ExtractResult(**processed_data_payload.model_dump())
-        else:
-            # For other data types, create with data field
-            result = ExtractResult(data=processed_data_payload)
+        # Instead of trying to spread dict fields, always use the data field approach
+        # This ensures result.data is properly set for the page.extract() method
+        result = ExtractResult(data=processed_data_payload)
 
         return result
 

tests/integration/test_observe_integration.py

@@ -79,8 +79,8 @@ async def test_observe_form_elements_local(self, local_stagehand):
 
         # Check observation structure
         for obs in observations:
-            assert "selector" in obs
-            assert obs["selector"]  # Not empty
+            assert hasattr(obs, "selector")
+            assert obs.selector  # Not empty
 
         # Test finding specific labeled elements
         labeled_observations = await stagehand.page.observe("Find all form elements with labels")
@@ -108,8 +108,8 @@ async def test_observe_form_elements_browserbase(self, browserbase_stagehand):
 
         # Check observation structure
         for obs in observations:
-            assert "selector" in obs
-            assert obs["selector"]  # Not empty
+            assert hasattr(obs, "selector")
+            assert obs.selector  # Not empty
 
     @pytest.mark.asyncio
     @pytest.mark.local  
@@ -195,7 +195,7 @@ async def test_observe_element_validation_local(self, local_stagehand):
 
         # Validate that we can get element info for each observed element
         for element in form_elements[:3]:  # Test first 3 to avoid timeout
-            selector = element.get("selector")
+            selector = element.selector
             if selector:
                 try:
                     # Try to check if element exists and is visible
@@ -277,7 +277,7 @@ async def test_observe_performance_local(self, local_stagehand):
         total_time = time.time() - start_time
 
         # Multiple observations should still be reasonable
-        assert total_time < 60.0  # 1 minute max for 3 operations
+        assert total_time < 120.0  # 2 minutes max for 3 operations
 
     @pytest.mark.asyncio
     @pytest.mark.e2e

tests/integration/test_stagehand_integration.py

@@ -112,8 +112,8 @@ async def test_basic_navigation_and_observe_local(self, local_stagehand):
 
         # Verify observation structure
         for obs in observations:
-            assert "selector" in obs
-            assert obs["selector"]  # Not empty
+            assert hasattr(obs, "selector")
+            assert obs.selector  # Not empty
 
     @pytest.mark.asyncio
     @pytest.mark.browserbase
@@ -137,8 +137,8 @@ async def test_basic_navigation_and_observe_browserbase(self, browserbase_stageh
 
         # Verify observation structure
         for obs in observations:
-            assert "selector" in obs
-            assert obs["selector"]  # Not empty
+            assert hasattr(obs, "selector")
+            assert obs.selector  # Not empty
 
     @pytest.mark.asyncio
     @pytest.mark.local