| """Unit tests for BrowserSession with CDP backend.""" |
|
|
| from __future__ import annotations |
|
|
| import asyncio |
| from unittest.mock import AsyncMock, MagicMock, patch |
|
|
| import pytest |
|
|
| from aistudio_api.infrastructure.browser.cdp_client import CDPPage |
| from aistudio_api.infrastructure.gateway.session import BrowserSession |
| from aistudio_api.infrastructure.gateway.wire_types import AistudioContent, AistudioPart |
|
|
|
|
| @pytest.fixture |
| def mock_cdp_page(): |
| page = MagicMock(spec=CDPPage) |
| page.url = "https://aistudio.google.com/prompts/new_chat" |
| page.is_closed = MagicMock(return_value=False) |
| page.goto = AsyncMock() |
| page.evaluate = AsyncMock() |
| page.fill = AsyncMock() |
| page.click = AsyncMock() |
| page.wait_for_timeout = AsyncMock() |
| page.wait_for_selector = AsyncMock(return_value=True) |
| page.query_selector = AsyncMock(return_value=True) |
| page.send_control_enter = AsyncMock(return_value=True) |
| page.get_cookies = AsyncMock(return_value=[]) |
| page.set_cookies = AsyncMock() |
| page.on_request = MagicMock(return_value=lambda: None) |
| page.on_response = MagicMock(return_value=lambda: None) |
| return page |
|
|
|
|
| @pytest.mark.anyio |
| async def test_browser_session_generate_snapshot(mock_cdp_page): |
| session = BrowserSession(port=9222) |
| session._page = mock_cdp_page |
| session._snap_key = "test_snapshot_fn" |
|
|
| async def fake_evaluate(expr, *args, **kwargs): |
| if "Promise.resolve(result)" in expr or "Promise.resolve(dms[snapKey]" in expr: |
| return "!mocked_snapshot_token_value_123" |
| if "window.__bg_hooked" in expr or "return 'already_hooked'" in expr: |
| return "already_hooked" |
| if "!window.__bg_service" in expr: |
| return True |
| if "window.__sl" in expr: |
| return 100 |
| if "window.__sr" in expr: |
| return "!mocked_snapshot_token_value_123" |
| return None |
| mock_cdp_page.evaluate.side_effect = fake_evaluate |
|
|
| contents = [ |
| AistudioContent( |
| role="user", |
| parts=[AistudioPart(text="hello world")], |
| ) |
| ] |
| snapshot = await session.generate_snapshot(contents) |
| assert snapshot == "!mocked_snapshot_token_value_123" |
|
|
|
|
| @pytest.mark.anyio |
| async def test_browser_session_send_hooked_request(mock_cdp_page): |
| session = BrowserSession(port=9222) |
| session._page = mock_cdp_page |
| session._templates["test_model"] = { |
| "url": "https://alkalimakersuite-pa.clients6.google.com/test", |
| "headers": {"content-type": "application/json"}, |
| } |
|
|
| async def fake_evaluate(expr, args=None, *a, **kw): |
| if "window.default_MakerSuite" in expr or "window.__bg_hooked" in expr: |
| return "already_hooked" |
| if "!window.__bg_service" in expr: |
| return True |
| if "XMLHttpRequest" in expr: |
| return {"status": 200, "body": '["response_data"]'} |
| return None |
|
|
| mock_cdp_page.evaluate.side_effect = fake_evaluate |
|
|
| status, body = await session.send_hooked_request(body='["test"]', timeout_ms=5000) |
| assert status == 200 |
| assert body == b'["response_data"]' |
|
|
|
|
| @pytest.mark.anyio |
| async def test_browser_session_send_streaming_request(mock_cdp_page): |
| session = BrowserSession(port=9222) |
| session._page = mock_cdp_page |
| session._templates["test_model"] = { |
| "url": "https://alkalimakersuite-pa.clients6.google.com/stream", |
| "headers": {"content-type": "application/json"}, |
| } |
|
|
| events = [ |
| {"type": "status", "status": 200}, |
| {"type": "chunk", "text": "data: hello "}, |
| {"type": "chunk", "text": "world\n\n"}, |
| {"type": "done"}, |
| ] |
|
|
| async def fake_evaluate(expr, args=None, *a, **kw): |
| if "window.default_MakerSuite" in expr or "window.__bg_hooked" in expr: |
| return "already_hooked" |
| if "window.__bg_service" in expr: |
| return True |
| if "(rid) => window.__stream_next" in expr: |
| if events: |
| return events.pop(0) |
| return {"type": "done"} |
| if "STREAMING_INIT" in expr or "window.__streams" in expr: |
| return None |
| return None |
|
|
| mock_cdp_page.evaluate.side_effect = fake_evaluate |
|
|
| collected = [] |
| async for tag, data in session.send_streaming_request(body='["stream"]', timeout_ms=5000): |
| collected.append((tag, data)) |
|
|
| assert ("status", 200) in collected |
| assert ("chunk", b"data: hello ") in collected |
| assert ("chunk", b"world\n\n") in collected |
|
|
