Spaces:
Paused
Paused
File size: 10,203 Bytes
a5784e9 | 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 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 | """
Tests for browser_utils/models/startup.py
Covers the log cleanup and standardization changes:
- [State] tagged logs
- Condensed verbose messages
- Silent success pattern
"""
import json
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@pytest.fixture
def mock_page():
"""Create a mock Playwright page with proper sync/async methods."""
page = MagicMock()
page.url = "https://aistudio.google.com/prompts/new_chat"
page.evaluate = AsyncMock()
page.goto = AsyncMock()
# locator() is sync, returns sync Locator - use MagicMock
# but locator.first.inner_text() is async
return page
class TestSetModelFromPageDisplay:
"""Tests for _set_model_from_page_display function."""
@pytest.mark.asyncio
@patch("api_utils.server_state.state")
async def test_reads_model_from_page(self, mock_state_obj, mock_page):
"""Test reading model display from page."""
from browser_utils.models.startup import _set_model_from_page_display
mock_state_obj.current_ai_studio_model_id = None
mock_state_obj.parsed_model_list = []
mock_state_obj.model_list_fetch_event = None
# Mock the model name locator - correctly chain .first.inner_text
first_locator = MagicMock()
first_locator.inner_text = AsyncMock(return_value="gemini-2.0-flash")
model_locator = MagicMock()
model_locator.first = first_locator
mock_page.locator.return_value = model_locator
await _set_model_from_page_display(mock_page, set_storage=False)
assert mock_state_obj.current_ai_studio_model_id == "gemini-2.0-flash"
@pytest.mark.asyncio
@patch("api_utils.server_state.state")
async def test_no_log_when_model_unchanged(self, mock_state_obj, mock_page):
"""Test silent success when model ID is unchanged."""
from browser_utils.models.startup import _set_model_from_page_display
mock_state_obj.current_ai_studio_model_id = "gemini-2.0-flash"
mock_state_obj.parsed_model_list = []
mock_state_obj.model_list_fetch_event = None
first_locator = MagicMock()
first_locator.inner_text = AsyncMock(return_value="gemini-2.0-flash")
model_locator = MagicMock()
model_locator.first = first_locator
mock_page.locator.return_value = model_locator
await _set_model_from_page_display(mock_page, set_storage=False)
# Should remain unchanged
assert mock_state_obj.current_ai_studio_model_id == "gemini-2.0-flash"
@pytest.mark.asyncio
@patch("api_utils.server_state.state")
@patch("browser_utils.models.startup._verify_and_apply_ui_state")
async def test_set_storage_updates_local_storage(
self, mock_verify_ui, mock_state_obj, mock_page
):
"""Test localStorage update when set_storage=True."""
from browser_utils.models.startup import _set_model_from_page_display
mock_state_obj.current_ai_studio_model_id = None
mock_state_obj.parsed_model_list = []
mock_state_obj.model_list_fetch_event = None
mock_verify_ui.return_value = True
first_locator = MagicMock()
first_locator.inner_text = AsyncMock(return_value="gemini-pro")
model_locator = MagicMock()
model_locator.first = first_locator
mock_page.locator.return_value = model_locator
mock_page.evaluate.return_value = None # No existing localStorage
await _set_model_from_page_display(mock_page, set_storage=True)
# Verify localStorage was written
assert mock_page.evaluate.call_count >= 2 # read + write
class TestHandleInitialModelStateAndStorage:
"""Tests for _handle_initial_model_state_and_storage function."""
@pytest.mark.asyncio
@patch("api_utils.server_state.state")
@patch("browser_utils.models.startup._verify_ui_state_settings")
async def test_no_refresh_when_storage_valid(
self, mock_verify_ui, mock_state_obj, mock_page
):
"""Test no page refresh when localStorage is valid."""
from browser_utils.models.startup import _handle_initial_model_state_and_storage
mock_state_obj.current_ai_studio_model_id = None
mock_state_obj.parsed_model_list = []
mock_state_obj.model_list_fetch_event = None
# Valid localStorage with all required fields
valid_prefs = json.dumps(
{
"promptModel": "models/gemini-2.0-flash",
"isAdvancedOpen": True,
"areToolsOpen": True,
}
)
mock_page.evaluate.return_value = valid_prefs
mock_verify_ui.return_value = {"needsUpdate": False}
await _handle_initial_model_state_and_storage(mock_page)
# Model should be set from localStorage
assert mock_state_obj.current_ai_studio_model_id == "gemini-2.0-flash"
@pytest.mark.asyncio
@patch("api_utils.server_state.state")
async def test_refresh_when_storage_missing(self, mock_state_obj, mock_page):
"""Test page refresh when localStorage is missing."""
from browser_utils.models.startup import _handle_initial_model_state_and_storage
mock_state_obj.current_ai_studio_model_id = None
mock_state_obj.parsed_model_list = []
mock_state_obj.model_list_fetch_event = None
# No localStorage
mock_page.evaluate.return_value = None
# Mock the locator for model name - properly chain .first.inner_text
first_locator = MagicMock()
first_locator.inner_text = AsyncMock(return_value="gemini-pro")
model_locator = MagicMock()
model_locator.first = first_locator
mock_page.locator.return_value = model_locator
# Mock goto and navigation
mock_page.goto = AsyncMock()
with (
patch("browser_utils.models.startup.expect_async") as mock_expect,
patch(
"browser_utils.models.startup._verify_and_apply_ui_state"
) as mock_verify,
):
mock_expect.return_value.to_be_visible = AsyncMock()
mock_verify.return_value = True
await _handle_initial_model_state_and_storage(mock_page)
# Should have attempted page reload
assert mock_page.goto.called
@pytest.mark.asyncio
@patch("api_utils.server_state.state")
async def test_handles_json_decode_error(self, mock_state_obj, mock_page):
"""Test handling of invalid JSON in localStorage."""
from browser_utils.models.startup import _handle_initial_model_state_and_storage
mock_state_obj.current_ai_studio_model_id = None
mock_state_obj.parsed_model_list = []
mock_state_obj.model_list_fetch_event = None
# Invalid JSON
mock_page.evaluate.return_value = "invalid json {"
first_locator = MagicMock()
first_locator.inner_text = AsyncMock(return_value="gemini-pro")
model_locator = MagicMock()
model_locator.first = first_locator
mock_page.locator.return_value = model_locator
mock_page.goto = AsyncMock()
with (
patch("browser_utils.models.startup.expect_async") as mock_expect,
patch(
"browser_utils.models.startup._verify_and_apply_ui_state"
) as mock_verify,
):
mock_expect.return_value.to_be_visible = AsyncMock()
mock_verify.return_value = True
# Should handle error gracefully
await _handle_initial_model_state_and_storage(mock_page)
class TestStateTaggedLogging:
"""Tests to verify [State] tagged logging patterns."""
@pytest.mark.asyncio
@patch("api_utils.server_state.state")
@patch("browser_utils.models.startup.logger")
async def test_state_tag_on_refresh_needed(
self, mock_logger, mock_state_obj, mock_page
):
"""Verify [State] tag in log when refresh is needed."""
from browser_utils.models.startup import _handle_initial_model_state_and_storage
mock_state_obj.current_ai_studio_model_id = None
mock_state_obj.parsed_model_list = []
mock_state_obj.model_list_fetch_event = None
mock_page.evaluate.return_value = None # No localStorage
first_locator = MagicMock()
first_locator.inner_text = AsyncMock(return_value="gemini-pro")
model_locator = MagicMock()
model_locator.first = first_locator
mock_page.locator.return_value = model_locator
mock_page.goto = AsyncMock()
with (
patch("browser_utils.models.startup.expect_async") as mock_expect,
patch(
"browser_utils.models.startup._verify_and_apply_ui_state"
) as mock_verify,
):
mock_expect.return_value.to_be_visible = AsyncMock()
mock_verify.return_value = True
await _handle_initial_model_state_and_storage(mock_page)
# Verify [State] tagged debug log was called
debug_calls = [str(call) for call in mock_logger.debug.call_args_list]
assert any("[State]" in call for call in debug_calls)
@pytest.mark.asyncio
@patch("api_utils.server_state.state")
@patch("browser_utils.models.startup.logger")
async def test_model_tag_on_page_display(
self, mock_logger, mock_state_obj, mock_page
):
"""Verify [Model] tag when reading from page display."""
from browser_utils.models.startup import _set_model_from_page_display
mock_state_obj.current_ai_studio_model_id = None
mock_state_obj.parsed_model_list = []
mock_state_obj.model_list_fetch_event = None
first_locator = MagicMock()
first_locator.inner_text = AsyncMock(return_value="gemini-flash")
model_locator = MagicMock()
model_locator.first = first_locator
mock_page.locator.return_value = model_locator
await _set_model_from_page_display(mock_page, set_storage=False)
# Verify [Model] tagged debug log was called
debug_calls = [str(call) for call in mock_logger.debug.call_args_list]
assert any("[Model]" in call for call in debug_calls)
|