File size: 4,477 Bytes
95f8c4b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | """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
|