"""
Test that screenshots work correctly in headless browser mode.
"""
import asyncio
import base64
import time
import pytest
from browser_use.browser import BrowserProfile, BrowserSession
from browser_use.browser.events import NavigateToUrlEvent, ScreenshotEvent
class TestHeadlessScreenshots:
"""Test screenshot functionality specifically in headless browsers"""
@pytest.mark.skip(reason='TODO: fix')
async def test_screenshot_works_in_headless_mode(self, httpserver):
"""Explicitly test that screenshots can be captured in headless=True mode"""
# Create a browser session with headless=True
browser_session = BrowserSession(
browser_profile=BrowserProfile(
headless=True, # Explicitly set headless mode
user_data_dir=None,
keep_alive=False,
)
)
try:
# Start the session
await browser_session.start()
assert browser_session._cdp_client_root is not None
# Set up test page with visible content
httpserver.expect_request('/').respond_with_data(
"""
Headless Screenshot Test
This is a test page
Testing screenshot capture in headless mode
Red Box
""",
content_type='text/html',
)
# Navigate to test page
event = browser_session.event_bus.dispatch(NavigateToUrlEvent(url=httpserver.url_for('/')))
await event
await event.event_result(raise_if_any=True, raise_if_none=False)
# Take screenshot
screenshot_event = browser_session.event_bus.dispatch(ScreenshotEvent())
await screenshot_event
screenshot_b64 = await screenshot_event.event_result(raise_if_any=True, raise_if_none=False)
# Verify screenshot was captured
assert screenshot_b64 is not None
assert isinstance(screenshot_b64, str)
assert len(screenshot_b64) > 0
# Decode and validate the screenshot
screenshot_bytes = base64.b64decode(screenshot_b64)
# Verify PNG signature
assert screenshot_bytes.startswith(b'\x89PNG\r\n\x1a\n')
# Should be a reasonable size (not just a blank image)
assert len(screenshot_bytes) > 5000, f'Screenshot too small: {len(screenshot_bytes)} bytes'
# Test full page screenshot
screenshot_event = browser_session.event_bus.dispatch(ScreenshotEvent(full_page=True))
await screenshot_event
full_page_screenshot = await screenshot_event.event_result(raise_if_any=True, raise_if_none=False)
assert full_page_screenshot is not None
full_page_bytes = base64.b64decode(full_page_screenshot)
assert full_page_bytes.startswith(b'\x89PNG\r\n\x1a\n')
assert len(full_page_bytes) > 5000
finally:
await browser_session.kill()
@pytest.mark.skip(reason='TODO: fix')
async def test_screenshot_with_state_summary_in_headless(self, httpserver):
"""Test that get_state_summary includes screenshots in headless mode"""
browser_session = BrowserSession(
browser_profile=BrowserProfile(
headless=True,
user_data_dir=None,
keep_alive=False,
)
)
try:
await browser_session.start()
# Set up test page
httpserver.expect_request('/').respond_with_data(
'State Summary Test
',
content_type='text/html',
)
event = browser_session.event_bus.dispatch(NavigateToUrlEvent(url=httpserver.url_for('/')))
await event
await event.event_result(raise_if_any=True, raise_if_none=False)
# Get state summary
state = await browser_session.get_browser_state_summary(cache_clickable_elements_hashes=False)
# Verify screenshot is included
assert state.screenshot is not None
assert isinstance(state.screenshot, str)
assert len(state.screenshot) > 0
# Decode and validate
screenshot_bytes = base64.b64decode(state.screenshot)
assert screenshot_bytes.startswith(b'\x89PNG\r\n\x1a\n')
assert len(screenshot_bytes) > 1000
finally:
await browser_session.kill()
@pytest.mark.skip(reason='TODO: fix')
async def test_screenshot_graceful_handling_in_headless(self, httpserver):
"""Test that screenshot handling works correctly in headless mode even with closed pages"""
# Set up test page
httpserver.expect_request('/test').respond_with_data(
'Test Page
', content_type='text/html'
)
browser_session = BrowserSession(
browser_profile=BrowserProfile(
headless=True,
user_data_dir=None,
keep_alive=False,
)
)
try:
await browser_session.start()
# Skip complex page manipulation - CDP doesn't have direct pages access
pytest.skip('CDP pages access pattern needs refactoring')
# Browser should auto-create a new page on about:blank with animation
# With AboutBlankWatchdog, about:blank pages now have animated content, so they should have screenshots
state = await browser_session.get_browser_state_summary(cache_clickable_elements_hashes=False)
assert state.screenshot is not None, 'Screenshot should not be None for animated about:blank pages'
assert state.url == 'about:blank' or state.url.startswith('chrome://'), f'Expected empty page but got {state.url}'
# Now navigate to a real page and verify screenshot works
event = browser_session.event_bus.dispatch(NavigateToUrlEvent(url=httpserver.url_for('/test')))
await event
await event.event_result(raise_if_any=True, raise_if_none=False)
# Get state with screenshot
state = await browser_session.get_browser_state_summary(cache_clickable_elements_hashes=False)
# Should have a screenshot now
assert state.screenshot is not None, 'Screenshot should not be None for real pages'
assert isinstance(state.screenshot, str)
assert len(state.screenshot) > 100, 'Screenshot should have substantial content'
assert 'test' in state.url.lower()
finally:
await browser_session.kill()
@pytest.mark.skip(reason='TODO: fix')
async def test_parallel_screenshots_long_page(self, httpserver):
"""Test screenshots in a highly parallel environment with a very long page"""
# Generate a very long page (50,000px+)
long_content = []
long_content.append('Very Long Page')
long_content.append('')
# Add many div elements to create a 50,000px+ long page
# Each div is 500px tall, so we need 100+ divs
for i in range(120):
color = f'rgb({i % 256}, {(i * 2) % 256}, {(i * 3) % 256})'
long_content.append(
f''
f'Section {i + 1} - Testing Parallel Screenshots'
f'
'
)
long_content.append('')
html_content = ''.join(long_content)
# Set up the test page
httpserver.expect_request('/longpage').respond_with_data(
html_content,
content_type='text/html',
)
test_url = httpserver.url_for('/longpage')
# Create 10 browser sessions
browser_sessions = []
for i in range(10):
session = BrowserSession(
browser_profile=BrowserProfile(
headless=True,
user_data_dir=None,
keep_alive=False,
)
)
browser_sessions.append(session)
try:
# Start all sessions sequentially to avoid playwright_global_object semaphore contention
# The playwright global object semaphore only allows 1 concurrent initialization
print('Starting 10 browser sessions sequentially...')
for i, session in enumerate(browser_sessions):
print(f'Starting session {i + 1}/10...')
await session.start()
# Navigate all sessions to the long page in parallel
print('Navigating all sessions to the long test page...')
await asyncio.gather(*[session.navigate(test_url) for session in browser_sessions])
# Take screenshots from all sessions
# Due to semaphore_limit=1, these will execute sequentially
print('Taking screenshots from all 10 sessions...')
start_time = time.time()
screenshot_tasks = [session.take_screenshot() for session in browser_sessions]
# Use return_exceptions=True to handle any failures gracefully
results = await asyncio.gather(*screenshot_tasks, return_exceptions=True)
total_time = time.time() - start_time
# Verify timing - with semaphore_limit=1, screenshots execute sequentially
# Each screenshot should take ~1.5s, so 10 × 1.5s = 15s, allow up to 30s for overhead
assert total_time < 30, f'Screenshots took too long: {total_time:.1f}s (should be < 30s)'
print(f'All screenshot attempts completed in {total_time:.1f}s')
# Separate successful screenshots from failures
screenshots = []
failures = []
for i, result in enumerate(results):
if isinstance(result, Exception):
failures.append((i, result))
print(f'Session {i} failed: {type(result).__name__}: {result}')
else:
screenshots.append(result)
print(f'Session {i} screenshot completed successfully')
# ALL screenshots must succeed
assert len(failures) == 0, (
f'{len(failures)} screenshots failed: {[(i, type(e).__name__, str(e)) for i, e in failures]}'
)
assert len(screenshots) == 10, f'Expected 10 successful screenshots, got {len(screenshots)}'
print('✅ All 10 screenshots captured successfully!')
# Verify all screenshots are valid
print('Verifying all 10 screenshots...')
for i, screenshot in enumerate(screenshots):
# Should not be None
assert screenshot is not None, f'Screenshot {i} returned None'
assert isinstance(screenshot, str), f'Screenshot {i} is not a string'
assert len(screenshot) > 0, f'Screenshot {i} is empty'
# Decode and validate
try:
screenshot_bytes = base64.b64decode(screenshot)
except Exception as e:
raise AssertionError(f'Screenshot {i} is not valid base64: {e}')
# Verify PNG signature
assert screenshot_bytes.startswith(b'\x89PNG\r\n\x1a\n'), f'Screenshot {i} is not a valid PNG'
# Full page screenshot should be reasonably large
# Due to our 6,000px height limit, expect at least 5KB
assert len(screenshot_bytes) > 20, f'Screenshot {i} too small: {len(screenshot_bytes)} bytes'
if len(screenshot_bytes) < 500:
print(
f'⚠️ Screenshot {i} failed to be taken in time, it returned a blank image instead: {len(screenshot_bytes)} bytes, perhaps the page failed to load in time?'
)
print('✅ All 10 screenshots validated successfully!')
# Also test taking regular (viewport) screenshots
print('\nTaking viewport screenshots from all sessions...')
start_time = time.time()
viewport_results = await asyncio.gather(
*[session.take_screenshot() for session in browser_sessions], return_exceptions=True
)
viewport_time = time.time() - start_time
assert viewport_time < 30, f'Viewport screenshots took too long: {viewport_time:.1f}s (should be < 30s)'
print(f'All viewport screenshot attempts completed in {viewport_time:.1f}s')
# Check for failures
viewport_screenshots = []
viewport_failures = []
for i, result in enumerate(viewport_results):
if isinstance(result, Exception):
viewport_failures.append((i, result))
print(f'Session {i} viewport failed: {type(result).__name__}: {result}')
else:
viewport_screenshots.append(result)
print(f'Session {i} viewport screenshot completed successfully')
# ALL viewport screenshots must succeed
assert len(viewport_failures) == 0, (
f'{len(viewport_failures)} viewport screenshots failed: {[(i, type(e).__name__, str(e)) for i, e in viewport_failures]}'
)
assert len(viewport_screenshots) == 10, (
f'Expected 10 successful viewport screenshots, got {len(viewport_screenshots)}'
)
print('✅ All 10 viewport screenshots captured successfully!')
# Verify all 10 viewport screenshots
print('Verifying all 10 viewport screenshots...')
for i, screenshot in enumerate(viewport_screenshots):
assert screenshot is not None, f'Viewport screenshot {i} is None'
screenshot_bytes = base64.b64decode(screenshot)
assert screenshot_bytes.startswith(b'\x89PNG\r\n\x1a\n'), f'Viewport screenshot {i} is not a valid PNG'
# Viewport screenshots should be reasonably sized
assert len(screenshot_bytes) > 10, f'Viewport screenshot {i} too small: {len(screenshot_bytes)} bytes'
print('✅ All 10 viewport screenshots validated successfully!')
finally:
# Kill all sessions in parallel
print('Killing all browser sessions...')
# Use return_exceptions=True to prevent one failed kill from affecting others
# This prevents "Future exception was never retrieved" errors
results = await asyncio.gather(*[session.kill() for session in browser_sessions], return_exceptions=True)
# Check that no exceptions were raised during cleanup
for i, result in enumerate(results):
if isinstance(result, Exception):
print(f'Warning: Session {i} kill raised exception: {type(result).__name__}: {result}')
@pytest.mark.skip(reason='TODO: fix')
async def test_screenshot_at_bottom_of_page(self, httpserver):
"""Test screenshot capture when scrolled to bottom of page (regression test for clipping issue)"""
browser_session = BrowserSession(
browser_profile=BrowserProfile(
headless=True,
user_data_dir=None,
keep_alive=False,
)
)
try:
await browser_session.start()
# Create a page with scrollable content
httpserver.expect_request('/scrollable').respond_with_data(
"""
Scrollable Page Test
Top of page
Middle of page
Bottom of page
""",
content_type='text/html',
)
# Navigate to test page
event = browser_session.event_bus.dispatch(NavigateToUrlEvent(url=httpserver.url_for('/scrollable')))
await event
await event.event_result(raise_if_any=True, raise_if_none=False)
# Skip - get_current_tab doesn't exist in CDP session
pytest.skip('get_current_tab method not available in CDP session')
finally:
await browser_session.kill()
class TestScreenshotEventSystem:
"""Tests for NEW event-driven screenshot infrastructure only."""
@pytest.mark.skip(reason='TODO: fix')
async def test_screenshot_response_event_dispatching(self, httpserver):
"""Test that ScreenshotResponseEvent is properly dispatched by event handlers."""
from browser_use.browser.events import ScreenshotEvent
browser_session = BrowserSession(browser_profile=BrowserProfile(headless=True, user_data_dir=None, keep_alive=False))
try:
await browser_session.start()
# Set up test page
httpserver.expect_request('/event-test').respond_with_data(
'Screenshot Event Test
',
content_type='text/html',
)
event = browser_session.event_bus.dispatch(NavigateToUrlEvent(url=httpserver.url_for('/event-test')))
await event
await event.event_result(raise_if_any=True, raise_if_none=False)
# Test the NEW event-driven path: direct event dispatching
event = browser_session.event_bus.dispatch(ScreenshotEvent(full_page=False))
screenshot_b64 = await event.event_result()
assert screenshot_b64 is not None
assert isinstance(screenshot_b64, str)
assert len(base64.b64decode(screenshot_b64)) > 5000
finally:
await browser_session.kill()