From db6af39b4bfbe8a0b81c0256edf5a1c3cb3b617b Mon Sep 17 00:00:00 2001
From: Conor Branagan <conor.branagan@gmail.com>
Date: Sun, 22 Jun 2025 15:01:34 -0400
Subject: [PATCH] Fix call signature for local screenshot

This function takes a set of kwargs rather than a dictionary.
---
 stagehand/page.py                          |  2 +-
 tests/integration/local/test_core_local.py | 58 +++++++++++++++-------
 2 files changed, 41 insertions(+), 19 deletions(-)

diff --git a/stagehand/page.py b/stagehand/page.py
index 300d7b2..2b67d35 100644
--- a/stagehand/page.py
+++ b/stagehand/page.py
@@ -399,7 +399,7 @@ async def screenshot(self, options: Optional[dict] = None) -> str:
 
             return result
         else:
-            return await self._page.screenshot(options)
+            return await self._page.screenshot(**payload)
 
     # Method to get or initialize the persistent CDP client
     async def get_cdp_client(self) -> CDPSession:
diff --git a/tests/integration/local/test_core_local.py b/tests/integration/local/test_core_local.py
index 0eb11ab..e5a3fdb 100644
--- a/tests/integration/local/test_core_local.py
+++ b/tests/integration/local/test_core_local.py
@@ -45,34 +45,40 @@ async def test_stagehand_local_initialization(self, stagehand_local):
     async def test_local_observe_and_act_workflow(self, stagehand_local):
         """Test core observe and act workflow in LOCAL mode - extracted from e2e tests."""
         stagehand = stagehand_local
-        
+
         # Navigate to a form page for testing
         await stagehand.page.goto("https://httpbin.org/forms/post")
-        
+
         # Test OBSERVE primitive: Find form elements
         form_elements = await stagehand.page.observe("Find all form input elements")
-        
+
         # Verify observations
         assert form_elements is not None
         assert len(form_elements) > 0
-        
+
         # Verify observation structure
         for obs in form_elements:
             assert hasattr(obs, "selector")
             assert obs.selector  # Not empty
-        
+
         # Test ACT primitive: Fill form fields
-        await stagehand.page.act("Fill the customer name field with 'Local Integration Test'")
+        await stagehand.page.act(
+            "Fill the customer name field with 'Local Integration Test'"
+        )
         await stagehand.page.act("Fill the telephone field with '555-LOCAL'")
         await stagehand.page.act("Fill the email field with 'local@integration.test'")
-        
+
         # Verify actions worked by observing filled fields
-        filled_fields = await stagehand.page.observe("Find all filled form input fields")
+        filled_fields = await stagehand.page.observe(
+            "Find all filled form input fields"
+        )
         assert filled_fields is not None
         assert len(filled_fields) > 0
-        
+
         # Test interaction with specific elements
-        customer_field = await stagehand.page.observe("Find the customer name input field")
+        customer_field = await stagehand.page.observe(
+            "Find the customer name input field"
+        )
         assert customer_field is not None
         assert len(customer_field) > 0
 
@@ -82,17 +88,17 @@ async def test_local_observe_and_act_workflow(self, stagehand_local):
     async def test_local_basic_navigation_and_observe(self, stagehand_local):
         """Test basic navigation and observe functionality in LOCAL mode"""
         stagehand = stagehand_local
-        
+
         # Navigate to a simple page
         await stagehand.page.goto("https://example.com")
-        
+
         # Observe elements on the page
         observations = await stagehand.page.observe("Find all the links on the page")
-        
+
         # Verify we got some observations
         assert observations is not None
         assert len(observations) > 0
-        
+
         # Verify observation structure
         for obs in observations:
             assert hasattr(obs, "selector")
@@ -104,14 +110,30 @@ async def test_local_basic_navigation_and_observe(self, stagehand_local):
     async def test_local_extraction_functionality(self, stagehand_local):
         """Test extraction functionality in LOCAL mode"""
         stagehand = stagehand_local
-        
+
         # Navigate to a content-rich page
         await stagehand.page.goto("https://news.ycombinator.com")
-        
+
         # Extract article titles using simple string instruction
         articles_text = await stagehand.page.extract(
             "Extract the titles of the first 3 articles on the page as a JSON list"
         )
-        
+
         # Verify extraction worked
-        assert articles_text is not None 
\ No newline at end of file
+        assert articles_text is not None
+
+    @pytest.mark.asyncio
+    @pytest.mark.integration
+    @pytest.mark.local
+    async def test_local_screenshot_functionality(self, stagehand_local):
+        """Test screenshot functionality in LOCAL mode"""
+        stagehand = stagehand_local
+
+        # Navigate to a content-rich page
+        await stagehand.page.goto("https://news.ycombinator.com")
+
+        # Take a screenshot
+        screenshot = await stagehand.page.screenshot()
+
+        # Verify screenshot worked
+        assert screenshot is not None