172 lines
5.9 KiB
Python
172 lines
5.9 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from datetime import datetime, timedelta, timezone
|
|
from typing import TYPE_CHECKING, Any
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
from playwright.async_api import Browser, BrowserContext, Page, Playwright, async_playwright
|
|
|
|
from crawlee.browsers import PlaywrightBrowserController, PlaywrightPersistentBrowser
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import AsyncGenerator
|
|
from pathlib import Path
|
|
|
|
from yarl import URL
|
|
|
|
|
|
@pytest.fixture
|
|
async def playwright() -> AsyncGenerator[Playwright, None]:
|
|
async with async_playwright() as playwright:
|
|
yield playwright
|
|
|
|
|
|
@pytest.fixture
|
|
async def browser(playwright: Playwright) -> AsyncGenerator[Browser, None]:
|
|
browser = await playwright.chromium.launch()
|
|
yield browser
|
|
await browser.close()
|
|
|
|
|
|
@pytest.fixture
|
|
async def controller(browser: Browser) -> AsyncGenerator[PlaywrightBrowserController, None]:
|
|
controller = PlaywrightBrowserController(browser, max_open_pages_per_browser=2)
|
|
yield controller
|
|
await controller.close()
|
|
|
|
|
|
async def test_initial_state(browser: Browser) -> None:
|
|
controller = PlaywrightBrowserController(browser)
|
|
|
|
# Test initial state
|
|
assert controller.pages == []
|
|
assert controller.pages_count == 0
|
|
assert isinstance(controller.last_page_opened_at, datetime)
|
|
assert controller.idle_time < timedelta(seconds=1)
|
|
assert controller.has_free_capacity
|
|
|
|
|
|
@pytest.mark.run_alone
|
|
async def test_open_and_close_page(controller: PlaywrightBrowserController, server_url: URL) -> None:
|
|
page = await controller.new_page()
|
|
await page.goto(str(server_url))
|
|
|
|
assert page in controller.pages
|
|
assert controller.pages_count == 1
|
|
assert controller.last_page_opened_at <= datetime.now(timezone.utc)
|
|
|
|
await page.close()
|
|
|
|
assert page not in controller.pages
|
|
assert controller.pages_count == 0
|
|
|
|
|
|
async def test_max_open_pages_limit(controller: PlaywrightBrowserController) -> None:
|
|
page1 = await controller.new_page()
|
|
assert controller.pages_count == 1
|
|
|
|
page2 = await controller.new_page()
|
|
assert controller.pages_count == 2
|
|
|
|
with pytest.raises(ValueError, match=r'Cannot open more pages in this browser.'):
|
|
await controller.new_page()
|
|
|
|
assert controller.pages_count == 2
|
|
|
|
await page1.close()
|
|
assert controller.pages_count == 1
|
|
|
|
page3 = await controller.new_page()
|
|
assert controller.pages_count == 2
|
|
|
|
await page2.close()
|
|
await page3.close()
|
|
|
|
assert controller.pages == []
|
|
assert controller.pages_count == 0
|
|
|
|
|
|
async def test_idle_time(controller: PlaywrightBrowserController) -> None:
|
|
idle_time_before = controller.idle_time
|
|
await asyncio.sleep(1) # Simulate waiting
|
|
idle_time_after = controller.idle_time
|
|
assert idle_time_after > idle_time_before
|
|
|
|
|
|
async def test_close_browser_with_open_pages(browser: Browser) -> None:
|
|
controller = PlaywrightBrowserController(browser, max_open_pages_per_browser=2)
|
|
_ = await controller.new_page()
|
|
|
|
with pytest.raises(ValueError, match=r'Cannot close the browser while there are open pages.'):
|
|
await controller.close()
|
|
|
|
assert controller.pages_count == 1
|
|
assert controller.is_browser_connected
|
|
|
|
await controller.close(force=True)
|
|
|
|
assert controller.pages_count == 0
|
|
assert not controller.is_browser_connected
|
|
|
|
|
|
async def test_memory_leak_on_concurrent_context_creation() -> None:
|
|
"""Test that only one browser context is created when multiple pages are opened concurrently."""
|
|
|
|
# Prepare mocked browser with relevant methods and attributes
|
|
mocked_browser = AsyncMock()
|
|
mocked_context_launcher = AsyncMock()
|
|
mocked_context = AsyncMock(spec=BrowserContext)
|
|
|
|
mocked_context_launcher.return_value = mocked_context
|
|
mocked_context.new_page.return_value = AsyncMock(spec=Page)
|
|
|
|
async def delayed_launch_persistent_context(*args: Any, **kwargs: Any) -> Any:
|
|
"""Ensure that both calls to create context overlap in time."""
|
|
await asyncio.sleep(5) # Simulate delay in creation to make sure race condition happens
|
|
return await mocked_context_launcher(*args, **kwargs)
|
|
|
|
mocked_browser.launch_persistent_context = delayed_launch_persistent_context
|
|
|
|
# Create minimal instance of PlaywrightBrowserController with mocked browser
|
|
controller = PlaywrightBrowserController(
|
|
PlaywrightPersistentBrowser(mocked_browser, None, {}), header_generator=None, fingerprint_generator=None
|
|
)
|
|
|
|
# Both calls will try to create browser context at the same time, but only one context should be created.
|
|
await asyncio.gather(controller.new_page(), controller.new_page())
|
|
|
|
assert mocked_context_launcher.call_count == 1
|
|
|
|
|
|
async def test_max_open_pages_limit_on_concurrent_creation(controller: PlaywrightBrowserController) -> None:
|
|
pages = await asyncio.gather(controller.new_page(), controller.new_page())
|
|
|
|
assert controller.pages_count == 2
|
|
|
|
for page in pages:
|
|
await page.close()
|
|
|
|
|
|
async def test_max_open_pages_limit_error_on_concurrent_creation(controller: PlaywrightBrowserController) -> None:
|
|
"""Test that max open pages limit is respected during concurrent page creation."""
|
|
with pytest.raises(ValueError, match=r'Cannot open more pages in this browser.'):
|
|
await asyncio.gather(controller.new_page(), controller.new_page(), controller.new_page())
|
|
|
|
|
|
async def test_browser_with_pre_existing_context(tmp_path: Path) -> None:
|
|
"""Test that using `Browser` with pre-existing active context re-uses such context."""
|
|
async with async_playwright() as pw:
|
|
persistent_context = await pw.firefox.launch_persistent_context(
|
|
user_data_dir=str(tmp_path),
|
|
headless=True,
|
|
)
|
|
browser = persistent_context.browser
|
|
assert browser
|
|
|
|
controller = PlaywrightBrowserController(browser=browser)
|
|
page_1 = await controller.new_page()
|
|
page_2 = await controller.new_page()
|
|
assert page_1.context == page_2.context == persistent_context
|