didi-lot1-ai/ai_platform/modules/web/tests/test_browse_client.py

198 lines
7.3 KiB
Python

"""Unit tests for BrowseClient with mocked Playwright."""
import asyncio
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from web.config import WebSettings
@pytest.fixture
def settings() -> WebSettings:
"""Create test settings."""
return WebSettings(
searxng_base_url="http://localhost:55100",
llm_base_url="http://localhost:14011",
external_url="http://localhost:51100",
)
class TestBrowseClientImport:
"""Tests for BrowseClient import handling."""
def test_raises_import_error_without_playwright(
self, settings: WebSettings
) -> None:
"""Test that BrowseClient raises ImportError without playwright."""
with patch("web.browse.client.HAS_PLAYWRIGHT", False):
from web.browse.client import BrowseClient
with pytest.raises(ImportError, match="playwright"):
BrowseClient(settings)
class TestBrowseClientBrowse:
"""Tests for BrowseClient.browse with mocked browser."""
def _make_mock_page(self) -> MagicMock:
"""Create a fresh mock Playwright page."""
page = AsyncMock()
page.goto = AsyncMock()
page.title = AsyncMock(return_value="Test Page Title")
page.url = "https://example.com/page"
page.content = AsyncMock(return_value="<html><body>Content</body></html>")
page.evaluate = AsyncMock(
side_effect=[
None, # remove unwanted elements
"Main content of the page with sufficient text for testing.",
"https://example.com/canonical", # canonical URL
"2024-01-15T00:00:00Z", # publication date
]
)
page.screenshot = AsyncMock(return_value=b"\x89PNG\r\n")
return page
def _make_mock_context(self, mock_page: MagicMock | None = None) -> MagicMock:
"""Create a mock browser context."""
if mock_page is None:
mock_page = self._make_mock_page()
context = AsyncMock()
context.new_page = AsyncMock(return_value=mock_page)
context.close = AsyncMock()
return context
def _make_mock_browser(self, mock_context: MagicMock | None = None) -> MagicMock:
"""Create a mock browser that creates fresh contexts each time."""
if mock_context is not None:
browser = MagicMock()
browser.is_connected = MagicMock(return_value=True)
browser.new_context = AsyncMock(return_value=mock_context)
browser.close = AsyncMock()
return browser
# Create a browser that returns fresh page mocks each time
browser = MagicMock()
browser.is_connected = MagicMock(return_value=True)
async def _new_context(**kwargs):
return self._make_mock_context()
browser.new_context = AsyncMock(side_effect=_new_context)
browser.close = AsyncMock()
return browser
@pytest.mark.asyncio
async def test_browse_single_page(self, settings: WebSettings) -> None:
"""Test browsing a single page returns content."""
with patch("web.browse.client.HAS_PLAYWRIGHT", True):
from web.browse.client import BrowseClient
from web.schemas.browse import BrowseRequest
mock_browser = self._make_mock_browser()
client = BrowseClient(settings, browser=mock_browser)
request = BrowseRequest(urls=["https://example.com/page"])
response = await client.browse(request)
assert response.total_browsed == 1
assert response.pages[0].title == "Test Page Title"
assert "Main content" in response.pages[0].text
@pytest.mark.asyncio
async def test_browse_handles_timeout(self, settings: WebSettings) -> None:
"""Test that Playwright timeout is properly wrapped."""
class MockPlaywrightTimeout(Exception):
pass
with (
patch("web.browse.client.HAS_PLAYWRIGHT", True),
patch(
"web.browse.client.PlaywrightTimeout",
MockPlaywrightTimeout,
create=True,
),
):
from web.browse.client import BrowseClient
from web.schemas.browse import BrowseRequest
mock_page = AsyncMock()
mock_page.goto = AsyncMock(side_effect=MockPlaywrightTimeout("Timeout"))
mock_context = self._make_mock_context(mock_page)
mock_browser = self._make_mock_browser(mock_context)
client = BrowseClient(settings, browser=mock_browser)
request = BrowseRequest(urls=["https://slow.example.com"])
response = await client.browse(request)
assert response.total_failed == 1
assert response.total_browsed == 0
@pytest.mark.asyncio
async def test_browse_multiple_urls(self, settings: WebSettings) -> None:
"""Test browsing multiple URLs concurrently."""
with patch("web.browse.client.HAS_PLAYWRIGHT", True):
from web.browse.client import BrowseClient
from web.schemas.browse import BrowseRequest
# Use browser that creates fresh contexts/pages
mock_browser = self._make_mock_browser()
client = BrowseClient(settings, browser=mock_browser)
request = BrowseRequest(
urls=["https://example.com/a", "https://example.com/b"],
parallel_browses=2,
)
response = await client.browse(request)
assert response.total_browsed == 2
@pytest.mark.asyncio
async def test_close_only_closes_owned_browser(self, settings: WebSettings) -> None:
"""Test that close() only closes browser if client owns it."""
with patch("web.browse.client.HAS_PLAYWRIGHT", True):
from web.browse.client import BrowseClient
mock_browser = self._make_mock_browser()
client = BrowseClient(settings, browser=mock_browser)
await client.close()
mock_browser.close.assert_not_called()
@pytest.mark.asyncio
async def test_concurrent_get_browser_single_instance(
self, settings: WebSettings
) -> None:
"""Test that concurrent _get_browser calls produce a single browser."""
with patch("web.browse.client.HAS_PLAYWRIGHT", True):
from web.browse.client import BrowseClient
mock_browser = self._make_mock_browser()
launch_count = 0
async def mock_launch(**kwargs):
nonlocal launch_count
launch_count += 1
await asyncio.sleep(0.05) # Simulate async work
return mock_browser
mock_pw = AsyncMock()
mock_pw.chromium.launch = mock_launch
mock_pw_ctx = AsyncMock()
mock_pw_ctx.start = AsyncMock(return_value=mock_pw)
client = BrowseClient(settings)
with patch("web.browse.client.async_playwright", return_value=mock_pw_ctx):
# Launch many concurrent _get_browser calls
results = await asyncio.gather(
*[client._get_browser() for _ in range(5)]
)
# Lock ensures only one browser was launched
assert launch_count == 1
# All results should be the same instance
assert all(r is results[0] for r in results)