import asyncio import json import sys from unittest.mock import AsyncMock, MagicMock, PropertyMock, patch import pytest # Mock server module before importing model_management components if needed # But we already imported them. Let's patch where necessary. from browser_utils.model_management import ( _force_ui_state_settings, _force_ui_state_with_retry, _handle_initial_model_state_and_storage, _set_model_from_page_display, _verify_and_apply_ui_state, _verify_ui_state_settings, load_excluded_models, switch_ai_studio_model, ) @pytest.fixture def mock_page(): page = AsyncMock() # locator is synchronous in Playwright page.locator = MagicMock() # Default evaluate returns None (empty localStorage) page.evaluate.return_value = None return page @pytest.fixture def mock_server(): mock = MagicMock() mock.processing_lock = MagicMock() mock.processing_lock.__aenter__ = AsyncMock(return_value=None) mock.processing_lock.__aexit__ = AsyncMock(return_value=None) mock.excluded_model_ids = set() mock.current_ai_studio_model_id = "initial-model" mock.parsed_model_list = [] # Use AsyncMock for async event, not MagicMock mock.model_list_fetch_event = AsyncMock(spec=asyncio.Event) mock.model_list_fetch_event.is_set.return_value = True return mock @pytest.mark.asyncio @pytest.mark.timeout(5) @pytest.mark.parametrize( "evaluate_result,evaluate_side_effect,expected_exists,expected_needs_update,expected_advanced,expected_tools,expected_error_contains,test_id", [ ( json.dumps({"isAdvancedOpen": True, "areToolsOpen": True}), None, True, False, True, True, None, "valid", ), ( json.dumps({"isAdvancedOpen": False, "areToolsOpen": True}), None, True, True, None, None, None, "needs_update", ), (None, None, False, True, None, None, "localStorage not found", "missing"), ( "invalid-json", None, False, True, None, None, "JSON parse failed", "json_error", ), ( None, Exception("Eval Error"), False, True, None, None, "Verification failed", "eval_error", ), ], ) async def test_verify_ui_state_settings( mock_page, evaluate_result, evaluate_side_effect, expected_exists, expected_needs_update, expected_advanced, expected_tools, expected_error_contains, test_id, ): """Test UI state settings verification with various scenarios.""" if evaluate_side_effect: mock_page.evaluate.side_effect = evaluate_side_effect else: mock_page.evaluate.return_value = evaluate_result with patch("browser_utils.models.ui_state.logger"): result = await _verify_ui_state_settings(mock_page, "req1") assert result["exists"] is expected_exists assert result["needsUpdate"] is expected_needs_update if expected_advanced is not None: assert result["isAdvancedOpen"] is expected_advanced if expected_tools is not None: assert result["areToolsOpen"] is expected_tools if expected_error_contains: assert expected_error_contains in result["error"] @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_force_ui_state_settings_success(mock_page): # Initial state: needs update initial_prefs = {"isAdvancedOpen": False} with ( patch("browser_utils.models.ui_state._verify_ui_state_settings") as mock_verify, patch("browser_utils.models.ui_state.logger"), ): mock_verify.side_effect = [ {"needsUpdate": True, "prefs": initial_prefs}, # First call {"needsUpdate": False}, # Second call ] result = await _force_ui_state_settings(mock_page, "req1") assert result is True # Check if setItem was called assert mock_page.evaluate.call_count == 1 args = mock_page.evaluate.call_args[0] assert "localStorage.setItem" in args[0] # Check if prefs were updated to True saved_prefs = json.loads(args[1]) assert saved_prefs["isAdvancedOpen"] is True assert saved_prefs["areToolsOpen"] is True @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_force_ui_state_settings_no_update_needed(mock_page): with ( patch("browser_utils.models.ui_state._verify_ui_state_settings") as mock_verify, patch("browser_utils.models.ui_state.logger"), ): mock_verify.return_value = {"needsUpdate": False} result = await _force_ui_state_settings(mock_page, "req1") assert result is True mock_page.evaluate.assert_not_called() @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_force_ui_state_settings_fail_verify(mock_page): with ( patch("browser_utils.models.ui_state._verify_ui_state_settings") as mock_verify, patch("browser_utils.models.ui_state.logger"), ): mock_verify.side_effect = [ {"needsUpdate": True, "prefs": {}}, {"needsUpdate": True}, # Still needs update after set ] result = await _force_ui_state_settings(mock_page, "req1") assert result is False mock_page.evaluate.assert_called_once() @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_force_ui_state_with_retry_success(mock_page): with ( patch("browser_utils.models.ui_state._force_ui_state_settings") as mock_force, patch("browser_utils.models.ui_state.logger"), ): mock_force.side_effect = [False, True] # Fail first, succeed second result = await _force_ui_state_with_retry( mock_page, "req1", max_retries=3, retry_delay=0.01 ) assert result is True assert mock_force.call_count == 2 @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_force_ui_state_with_retry_fail(mock_page): with ( patch("browser_utils.models.ui_state._force_ui_state_settings") as mock_force, patch("browser_utils.models.ui_state.logger"), ): mock_force.return_value = False result = await _force_ui_state_with_retry( mock_page, "req1", max_retries=2, retry_delay=0.01 ) assert result is False assert mock_force.call_count == 2 @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_verify_and_apply_ui_state_needs_update(mock_page): with ( patch("browser_utils.models.ui_state._verify_ui_state_settings") as mock_verify, patch("browser_utils.models.ui_state._force_ui_state_with_retry") as mock_retry, patch("browser_utils.models.ui_state.logger"), ): mock_verify.return_value = { "exists": True, "isAdvancedOpen": False, "areToolsOpen": False, "needsUpdate": True, } mock_retry.return_value = True result = await _verify_and_apply_ui_state(mock_page, "req1") assert result is True mock_retry.assert_called_once() @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_verify_and_apply_ui_state_ok(mock_page): with ( patch("browser_utils.models.ui_state._verify_ui_state_settings") as mock_verify, patch("browser_utils.models.ui_state._force_ui_state_with_retry") as mock_retry, patch("browser_utils.models.ui_state.logger"), ): mock_verify.return_value = { "exists": True, "isAdvancedOpen": True, "areToolsOpen": True, "needsUpdate": False, } result = await _verify_and_apply_ui_state(mock_page, "req1") assert result is True mock_retry.assert_not_called() @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_load_excluded_models(tmp_path): # Create a dummy exclusion file d = tmp_path / "config" d.mkdir() p = d / "excluded_models.txt" p.write_text("model-a\nmodel-b\n", encoding="utf-8") # Mock server state mock_state = MagicMock() mock_state.excluded_model_ids = set() with ( patch("api_utils.server_state.state", mock_state), patch("os.path.exists") as mock_exists, patch("builtins.open", new_callable=MagicMock) as mock_open, patch("browser_utils.models.switcher.logger"), ): mock_exists.return_value = True mock_file = MagicMock() mock_file.__enter__.return_value = ["model-a\n", "model-b\n"] mock_open.return_value = mock_file load_excluded_models("excluded_models.txt") assert "model-a" in mock_state.excluded_model_ids assert "model-b" in mock_state.excluded_model_ids @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_switch_ai_studio_model_already_set(mock_page): model_id = "gemini-pro" full_model_path = f"models/{model_id}" prefs = {"promptModel": full_model_path} mock_page.evaluate.return_value = json.dumps(prefs) mock_page.url = "https://aistudio.google.com/prompts/new_chat" with ( patch("browser_utils.models.switcher.logger"), patch("browser_utils.models.switcher.expect_async") as mock_expect, ): mock_expect.return_value.to_be_visible = AsyncMock() result = await switch_ai_studio_model(mock_page, model_id, "req1") assert result is True # Should not navigate if already on new_chat mock_page.goto.assert_not_called() @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_switch_ai_studio_model_success(mock_page): model_id = "gemini-pro" full_model_path = f"models/{model_id}" initial_prefs = {"promptModel": "models/other-model"} # Mock server module mock_server = MagicMock() mock_server.parsed_model_list = [{"id": model_id, "display_name": "Gemini Pro"}] with ( patch.dict(sys.modules, {"server": mock_server}), patch( "browser_utils.models.switcher._verify_and_apply_ui_state", return_value=True, ), patch("browser_utils.models.switcher.logger"), patch("browser_utils.models.switcher.expect_async") as mock_expect, ): mock_expect.return_value.to_be_visible = AsyncMock() # Stateful evaluate mock evaluate_mock = AsyncMock() mock_page.evaluate = evaluate_mock call_count = 0 async def evaluate_side_effect(script, *args): nonlocal call_count if "localStorage.getItem" in script: call_count += 1 if call_count == 1: # Initial check return json.dumps(initial_prefs) if call_count == 2: # Final verification return json.dumps({"promptModel": full_model_path}) return None evaluate_mock.side_effect = evaluate_side_effect # Mock page elements mock_locator = MagicMock() mock_locator.first.inner_text = AsyncMock( return_value=model_id ) # Matches target mock_page.locator.return_value = mock_locator # Mock incognito button mock_incognito = MagicMock() mock_incognito.get_attribute = AsyncMock(return_value="ms-button-active") def locator_side_effect(selector): if 'data-test-id="model-name"' in selector: return mock_locator if "Temporary chat toggle" in selector: return mock_incognito return MagicMock() mock_page.locator.side_effect = locator_side_effect result = await switch_ai_studio_model(mock_page, model_id, "req1") assert result is True mock_page.goto.assert_called() @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_set_model_from_page_display(mock_page): # Mock server state mock_state = MagicMock() mock_state.current_ai_studio_model_id = "old-model" # CRITICAL: Use AsyncMock for async event, not MagicMock mock_event = AsyncMock(spec=asyncio.Event) mock_event.is_set.return_value = True # Event is already set, no wait needed mock_state.model_list_fetch_event = mock_event mock_state.parsed_model_list = [] # Mock locator mock_locator = MagicMock() mock_locator.first.inner_text = AsyncMock(return_value="new-model") mock_page.locator.return_value = mock_locator with ( patch("api_utils.server_state.state", mock_state), patch("browser_utils.models.startup.logger"), ): await _set_model_from_page_display(mock_page, set_storage=False) assert mock_state.current_ai_studio_model_id == "new-model" @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_handle_initial_model_state_needs_reload(mock_page): # Mock empty localStorage -> needs reload mock_page.evaluate.return_value = None mock_page.url = "http://test.url" # Mock server mock_server = MagicMock() with ( patch.dict(sys.modules, {"server": mock_server}), patch( "browser_utils.models.startup._set_model_from_page_display" ) as mock_set_model, patch( "browser_utils.models.startup._verify_and_apply_ui_state", return_value=True, ), patch("browser_utils.models.startup.logger"), patch("browser_utils.models.startup.expect_async") as mock_expect, ): mock_expect.return_value.to_be_visible = AsyncMock() await _handle_initial_model_state_and_storage(mock_page) # Should call _set_model_from_page_display twice assert mock_set_model.call_count == 2 assert mock_set_model.call_args_list[0][1]["set_storage"] is True assert mock_set_model.call_args_list[1][1]["set_storage"] is False # Should reload page mock_page.goto.assert_called_with( "http://test.url", wait_until="domcontentloaded", timeout=40000 ) @pytest.mark.asyncio async def test_switch_ai_studio_model_revert_logic(mock_page): """Test the revert logic when model switch fails validation.""" model_id = "gemini-pro" full_model_path = f"models/{model_id}" initial_prefs = {"promptModel": "models/original-model"} original_prefs_str = json.dumps(initial_prefs) # Mock server module mock_server = MagicMock() mock_server.parsed_model_list = [{"id": model_id, "display_name": "Gemini Pro"}] with ( patch.dict(sys.modules, {"server": mock_server}), patch( "browser_utils.models.switcher._verify_and_apply_ui_state", return_value=True, ), patch("browser_utils.models.switcher.logger"), patch("browser_utils.models.switcher.expect_async") as mock_expect, ): mock_expect.return_value.to_be_visible = AsyncMock() # Setup stateful evaluate mock to simulate storage state changes call_count = 0 async def evaluate_side_effect(script, *args): nonlocal call_count if "localStorage.getItem" in script: call_count += 1 if call_count == 1: # Initial check return original_prefs_str elif call_count == 2: # Final verification (switched) return json.dumps({"promptModel": full_model_path}) elif call_count == 3: # Revert logic return json.dumps({"promptModel": full_model_path}) return None mock_page.evaluate = AsyncMock(side_effect=evaluate_side_effect) # Simulate mismatch: page displays "Original Model" but storage has new model mock_locator = MagicMock() mock_locator.first.inner_text = AsyncMock(return_value="Original Model") mock_page.locator.return_value = mock_locator # Execute - should fail validation result = await switch_ai_studio_model(mock_page, model_id, "req1") # Verify failure occurred assert result is False # Check setItem calls set_item_calls = [ args for args in mock_page.evaluate.call_args_list if "localStorage.setItem" in args[0][0] ] # Verify target prefs were set, but no revert assert len(set_item_calls) == 2 last_set_call = set_item_calls[-1] target_prefs = json.loads(last_set_call[0][1]) assert target_prefs["promptModel"] == full_model_path assert target_prefs["isAdvancedOpen"] is True assert target_prefs["areToolsOpen"] is True @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_switch_ai_studio_model_incognito_toggle(mock_page): """Test incognito toggle logic when model switch succeeds""" model_id = "gemini-pro" full_model_path = f"models/{model_id}" # Mock server module mock_server = MagicMock() mock_server.parsed_model_list = [{"id": model_id, "display_name": "Gemini Pro"}] with ( patch.dict(sys.modules, {"server": mock_server}), patch( "browser_utils.models.switcher._verify_and_apply_ui_state", return_value=True, ), patch("browser_utils.models.switcher.logger"), patch("browser_utils.models.switcher.expect_async") as mock_expect, ): mock_expect.return_value.to_be_visible = AsyncMock() # Mock evaluate for success path # It needs to return a DIFFERENT model initially so it doesn't return early call_count = 0 def evaluate_side_effect(script, *args): nonlocal call_count if "localStorage.getItem" in script: call_count += 1 if call_count == 1: # Initial check -> return old model return json.dumps({"promptModel": "models/old-model"}) if call_count == 2: # Final verification -> return new model return json.dumps({"promptModel": full_model_path}) return None mock_page.evaluate.side_effect = evaluate_side_effect # Mock page elements mock_locator = MagicMock() mock_locator.first.inner_text = AsyncMock(return_value=model_id) # Mock incognito button - INACTIVE initially mock_incognito = MagicMock() mock_incognito.wait_for = AsyncMock() mock_incognito.click = AsyncMock() # First check returns inactive, second check (after click) returns active mock_incognito.get_attribute = AsyncMock( side_effect=[ "ms-button", # inactive "ms-button-active ms-button", # active ] ) def locator_side_effect(selector): if 'data-test-id="model-name"' in selector: return mock_locator if "Temporary chat toggle" in selector: return mock_incognito return MagicMock() mock_page.locator.side_effect = locator_side_effect result = await switch_ai_studio_model(mock_page, model_id, "req1") assert result is True mock_incognito.click.assert_called_once() @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_exception_handling_coverage(mock_page): """Cover exception handlers in various functions""" # 1. _force_ui_state_settings exception with ( patch( "browser_utils.models.ui_state._verify_ui_state_settings", side_effect=Exception("Force Error"), ), patch("browser_utils.models.ui_state.logger"), ): assert await _force_ui_state_settings(mock_page) is False # 2. _verify_and_apply_ui_state exception with ( patch( "browser_utils.models.ui_state._verify_ui_state_settings", side_effect=Exception("Verify Apply Error"), ), patch("browser_utils.models.ui_state.logger"), ): assert await _verify_and_apply_ui_state(mock_page) is False # 3. switch_ai_studio_model JSON decode error is tested in test_handle_initial_model_state_exceptions # The JSONDecodeError handler falls back to empty prefs and continues execution @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_switch_ai_studio_model_nav_only(mock_page): """Test navigation when model already matches but URL is wrong""" model_id = "gemini-pro" full_model_path = f"models/{model_id}" prefs = {"promptModel": full_model_path} mock_page.evaluate.return_value = json.dumps(prefs) mock_page.url = "https://other.url" # Not new_chat with ( patch("browser_utils.models.switcher.logger"), patch("browser_utils.models.switcher.expect_async") as mock_expect, ): mock_expect.return_value.to_be_visible = AsyncMock() result = await switch_ai_studio_model(mock_page, model_id, "req1") assert result is True # Should navigate mock_page.goto.assert_called() @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_load_excluded_models_edge_cases(tmp_path): """Test edge cases for load_excluded_models""" # 1. File does not exist # Mock server module - use api_utils.server_state.state which is what the implementation uses mock_state = MagicMock() mock_state.excluded_model_ids = set() with ( patch("api_utils.server_state.state", mock_state), patch("browser_utils.models.switcher.logger") as mock_logger, ): load_excluded_models("non_existent.txt") # Implementation uses logger.debug, not logger.info debug_calls = [call[0][0] for call in mock_logger.debug.call_args_list] assert any("not found" in msg.lower() for msg in debug_calls) # 2. File exists but is empty - tested in the next block with mocked file I/O # Let's mock os.path.exists/open for easier testing of logic with ( patch("api_utils.server_state.state", mock_state), patch("os.path.exists", return_value=True), patch("builtins.open", new_callable=MagicMock) as mock_open, patch("browser_utils.models.switcher.logger") as mock_logger, ): # Empty file mock_file = MagicMock() mock_file.__enter__.return_value = [] # Empty lines mock_open.return_value = mock_file load_excluded_models("empty.txt") # Implementation uses logger.debug debug_calls = [call[0][0] for call in mock_logger.debug.call_args_list] assert any("empty" in msg.lower() for msg in debug_calls) # 3. Exception with ( patch("api_utils.server_state.state", mock_state), patch("os.path.exists", side_effect=Exception("Disk Error")), patch("browser_utils.models.switcher.logger") as mock_logger, ): load_excluded_models("error.txt") assert mock_logger.error.called @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_handle_initial_model_state_exceptions(mock_page): """Test exception handling in _handle_initial_model_state_and_storage""" mock_server = MagicMock() # 1. JSON Decode Error mock_page.evaluate.return_value = "invalid-json" with ( patch.dict(sys.modules, {"server": mock_server}), patch( "browser_utils.models.startup._set_model_from_page_display" ) as mock_set_model, patch("browser_utils.models.startup.logger") as mock_logger, ): # Should trigger reload path due to JSON error # We'll mock _set_model_from_page_display to raise Exception to test the outer try/except mock_set_model.side_effect = Exception("Inner Error") await _handle_initial_model_state_and_storage(mock_page) # Verify error log # Check that we have the catastrophic error log # It catches "Inner Error" in the outer except block error_calls = [args[0][0] for args in mock_logger.error.call_args_list] assert any("critical error" in msg.lower() for msg in error_calls) @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_handle_initial_model_state_reload_retry_logic(mock_page): """Test reload retry logic""" mock_server = MagicMock() mock_page.evaluate.return_value = None # Trigger reload with ( patch.dict(sys.modules, {"server": mock_server}), patch( "browser_utils.models.startup._set_model_from_page_display" ) as mock_set_model, patch("browser_utils.models.startup.logger"), patch("asyncio.sleep", new_callable=AsyncMock), patch("browser_utils.models.startup.expect_async") as mock_expect, ): mock_expect.return_value.to_be_visible = AsyncMock() # Mock goto to fail twice then succeed mock_page.goto.side_effect = [Exception("Fail 1"), Exception("Fail 2"), None] await _handle_initial_model_state_and_storage(mock_page) assert mock_page.goto.call_count == 3 # Should eventually succeed and call set_model twice (start and end) assert mock_set_model.call_count == 2 @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_set_model_from_page_display_timeout(mock_page): """Test timeout when waiting for model list""" mock_state = MagicMock() mock_state.current_ai_studio_model_id = None mock_state.parsed_model_list = [] # Use AsyncMock for async event mock_event = AsyncMock(spec=asyncio.Event) mock_event.is_set.return_value = False # Event not set, will wait mock_state.model_list_fetch_event = mock_event # Mock locator mock_locator = MagicMock() mock_locator.first.inner_text = AsyncMock(return_value="displayed-model") mock_page.locator.return_value = mock_locator with ( patch("api_utils.server_state.state", mock_state), patch("browser_utils.models.startup.logger") as mock_logger, patch("asyncio.wait_for", side_effect=asyncio.TimeoutError), ): await _set_model_from_page_display(mock_page, set_storage=False) # Should log warning about timeout assert any( "timeout waiting for model list" in str(arg).lower() for arg in mock_logger.warning.call_args_list[0][0] ) # Should still update global ID using display name as fallback assert mock_state.current_ai_studio_model_id == "displayed-model" @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_set_model_from_page_display_storage_logic(mock_page): """Test storage update logic in _set_model_from_page_display""" mock_state = MagicMock() mock_state.current_ai_studio_model_id = "old" # Create a set event to avoid hanging on wait() mock_event = asyncio.Event() mock_event.set() mock_state.model_list_fetch_event = mock_event mock_state.parsed_model_list = [] mock_locator = MagicMock() mock_locator.first.inner_text = AsyncMock(return_value="new-model") mock_page.locator.return_value = mock_locator # Mock existing prefs existing_prefs = {"someKey": "someVal"} mock_page.evaluate.return_value = json.dumps(existing_prefs) with ( patch("api_utils.server_state.state", mock_state), patch( "browser_utils.models.startup._verify_and_apply_ui_state", return_value=True, ), patch("browser_utils.models.startup.logger"), ): await _set_model_from_page_display(mock_page, set_storage=True) # Check that setItem was called with updated prefs assert mock_page.evaluate.call_count == 2 # getItem, setItem # Verify setItem args set_call = mock_page.evaluate.call_args_list[1] assert "localStorage.setItem" in set_call[0][0] saved_prefs = json.loads(set_call[0][1]) assert saved_prefs["isAdvancedOpen"] is True assert saved_prefs["promptModel"] == "models/new-model" # Check default keys added assert "bidiModel" in saved_prefs @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_switch_ai_studio_model_catastrophic_error(mock_page): """Test top-level exception handling in switch_ai_studio_model""" # Force an error immediately mock_page.evaluate.side_effect = Exception("Catastrophic Failure") with ( patch("browser_utils.models.switcher.logger") as mock_logger, patch( "browser_utils.operations.save_error_snapshot", new_callable=AsyncMock ) as mock_snapshot, ): result = await switch_ai_studio_model(mock_page, "model-id", "req1") assert result is False mock_snapshot.assert_called() assert mock_logger.exception.called # ============================================================================ # EXTENDED COVERAGE - Edge Cases, Error Handling, Revert Logic # ============================================================================ # === Section 1: UI State Verification Tests === @pytest.mark.asyncio async def test_verify_ui_state_missing_storage(mock_page): """Test verification when localStorage item is missing.""" mock_page.evaluate.return_value = None result = await _verify_ui_state_settings(mock_page) assert result["exists"] is False assert result["error"] == "localStorage not found" assert result["needsUpdate"] is True @pytest.mark.asyncio async def test_verify_ui_state_json_error(mock_page): """Test verification with invalid JSON in storage.""" mock_page.evaluate.return_value = "invalid json" result = await _verify_ui_state_settings(mock_page) assert result["exists"] is False assert "JSON parse failed" in result["error"] assert result["needsUpdate"] is True @pytest.mark.asyncio async def test_verify_ui_state_exception(mock_page): """Test verification when evaluate raises exception.""" mock_page.evaluate.side_effect = Exception("Page error") result = await _verify_ui_state_settings(mock_page) assert result["exists"] is False assert "Verification failed" in result["error"] assert result["needsUpdate"] is True # === Section 2: Force UI State Tests === @pytest.mark.asyncio async def test_force_ui_state_already_correct(mock_page): """Test force update when state is already correct.""" mock_page.evaluate.return_value = json.dumps( {"isAdvancedOpen": True, "areToolsOpen": True} ) result = await _force_ui_state_settings(mock_page) assert result is True assert mock_page.evaluate.call_count == 1 @pytest.mark.asyncio async def test_force_ui_state_verify_failure(mock_page): """Test force update when verification fails after setting.""" initial_state = json.dumps({"isAdvancedOpen": False}) final_state = json.dumps({"isAdvancedOpen": False}) # Still wrong mock_page.evaluate.side_effect = [ initial_state, # verify 1 None, # setItem final_state, # verify 2 ] result = await _force_ui_state_settings(mock_page) assert result is False @pytest.mark.asyncio async def test_force_ui_state_exception(mock_page): """Test force update exception handling.""" mock_page.evaluate.side_effect = Exception("Set error") result = await _force_ui_state_settings(mock_page) assert result is False @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_force_ui_state_settings_exception_via_verify(mock_page): """Test exception handling in _force_ui_state_settings via verify.""" with patch( "browser_utils.models.ui_state._verify_ui_state_settings", side_effect=Exception("Test Error"), ): result = await _force_ui_state_settings(mock_page, "req1") assert result is False @pytest.mark.asyncio @pytest.mark.parametrize( "mock_force_side_effect,mock_force_return_value,max_retries,retry_delay,expected_result,expected_call_count,test_id", [ ([False, True], None, 3, None, True, 2, "success"), (None, False, 2, 0.01, False, 2, "failure"), ], ) async def test_force_ui_state_retry( mock_page, mock_force_side_effect, mock_force_return_value, max_retries, retry_delay, expected_result, expected_call_count, test_id, ): """Test retry logic for forcing UI state settings (eventually succeeding or failing all attempts).""" with patch("browser_utils.models.ui_state._force_ui_state_settings") as mock_force: if mock_force_side_effect: mock_force.side_effect = mock_force_side_effect else: mock_force.return_value = mock_force_return_value kwargs = {"max_retries": max_retries} if retry_delay is not None: kwargs["retry_delay"] = retry_delay result = await _force_ui_state_with_retry(mock_page, **kwargs) assert result is expected_result assert mock_force.call_count == expected_call_count @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_verify_and_apply_ui_state_exception(mock_page): """Test exception handling in _verify_and_apply_ui_state.""" with patch( "browser_utils.models.ui_state._verify_ui_state_settings", side_effect=Exception("Test Error"), ): result = await _verify_and_apply_ui_state(mock_page, "req1") assert result is False # === Section 3: Load Excluded Models Tests === @pytest.mark.asyncio async def test_load_excluded_models_success(mock_server): """Test loading excluded models from file.""" mock_content = "model-1\nmodel-2\n" mock_file = MagicMock() mock_file.__enter__.return_value = mock_content.splitlines() with ( patch("builtins.open", return_value=mock_file), patch("os.path.exists", return_value=True), patch.dict(sys.modules, {"server": mock_server}), ): load_excluded_models("excluded.txt") assert "model-1" in mock_server.excluded_model_ids assert "model-2" in mock_server.excluded_model_ids # === Section 4: Handle Initial Model State Tests === @pytest.mark.asyncio async def test_handle_initial_state_missing_storage(mock_page, mock_server): """Test handling initial state when storage is missing.""" mock_page.evaluate.return_value = None # No storage mock_page.url = "http://test.url" with ( patch( "browser_utils.models.startup._set_model_from_page_display" ) as mock_set_model, patch.dict(sys.modules, {"server": mock_server}), patch( "browser_utils.models.startup._verify_and_apply_ui_state", return_value=True, ), patch("browser_utils.models.startup.expect_async") as mock_expect_async, ): mock_expect = MagicMock() mock_expect.to_be_visible = AsyncMock() mock_expect_async.return_value = mock_expect await _handle_initial_model_state_and_storage(mock_page) assert mock_set_model.call_count == 2 assert mock_page.goto.call_count == 1 @pytest.mark.asyncio async def test_handle_initial_state_valid_no_reload(mock_page, mock_server): """Test handling initial state when everything is valid.""" mock_page.evaluate.return_value = json.dumps( {"promptModel": "models/valid-model", "isAdvancedOpen": True} ) with patch("browser_utils.models.startup._verify_ui_state_settings") as mock_verify: mock_verify.return_value = {"needsUpdate": False} with patch.dict(sys.modules, {"server": mock_server}): await _handle_initial_model_state_and_storage(mock_page) mock_page.goto.assert_not_called() assert mock_server.current_ai_studio_model_id == "valid-model" @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_handle_initial_model_state_and_storage_success(mock_page, mock_server): """Test successful initial model state handling.""" initial_prefs = json.dumps( { "promptModel": "models/gemini-pro", "isAdvancedOpen": True, "areToolsOpen": True, } ) mock_page.evaluate.return_value = initial_prefs with ( patch( "browser_utils.models.startup._verify_ui_state_settings", return_value={"needsUpdate": False}, ), patch.dict("sys.modules", {"server": mock_server}), ): await _handle_initial_model_state_and_storage(mock_page) assert mock_server.current_ai_studio_model_id == "gemini-pro" mock_page.goto.assert_not_called() @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_handle_initial_model_state_exception(mock_page): """Test exception handling in _handle_initial_model_state_and_storage.""" mock_page.evaluate.side_effect = Exception("Init Error") with patch( "browser_utils.models.startup._set_model_from_page_display" ) as mock_fallback: await _handle_initial_model_state_and_storage(mock_page) mock_fallback.assert_called_once() @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_handle_initial_model_state_json_error(mock_page, mock_server): """Test JSON error handling in initial state check.""" mock_page.evaluate.return_value = "invalid-json" with ( patch.dict(sys.modules, {"server": mock_server}), patch("browser_utils.models.startup.logger") as mock_logger, patch( "browser_utils.models.startup._set_model_from_page_display" ) as mock_set_model, patch("browser_utils.models.startup.expect_async") as mock_expect, patch( "browser_utils.models.startup._verify_and_apply_ui_state", return_value=True, ), ): mock_expect.return_value.to_be_visible = AsyncMock() await _handle_initial_model_state_and_storage(mock_page) mock_set_model.assert_called() errors = [call.args[0] for call in mock_logger.error.call_args_list] assert any( "Failed to parse localStorage.aiStudioUserPreference JSON" in e for e in errors ) @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_handle_initial_model_state_reload_retry(mock_page, mock_server): """Test reload retry logic.""" mock_server = MagicMock() mock_page.evaluate.return_value = None # Trigger reload mock_page.goto.side_effect = [Exception("Load failed"), None] with ( patch.dict(sys.modules, {"server": mock_server}), patch("browser_utils.models.startup.logger") as mock_logger, patch("browser_utils.models.startup._set_model_from_page_display"), patch("browser_utils.models.startup.expect_async") as mock_expect, patch( "browser_utils.models.startup._verify_and_apply_ui_state", return_value=True, ), patch("asyncio.sleep", new_callable=AsyncMock), ): mock_expect.return_value.to_be_visible = AsyncMock() await _handle_initial_model_state_and_storage(mock_page) assert mock_page.goto.call_count == 2 warnings = [call.args[0] for call in mock_logger.warning.call_args_list] assert any("page reload attempt 1/3 failed" in w.lower() for w in warnings) # === Section 5: Set Model from Page Display Tests === @pytest.mark.asyncio async def test_set_model_from_display_basic(mock_page, mock_server): """Test setting model from page display.""" mock_element = AsyncMock() mock_element.inner_text.return_value = "displayed-model" mock_locator = AsyncMock() mock_locator.first = mock_element mock_page.locator.return_value = mock_locator mock_server.current_ai_studio_model_id = "old-model" with patch.dict(sys.modules, {"server": mock_server}): await _set_model_from_page_display(mock_page, set_storage=False) assert mock_server.current_ai_studio_model_id == "displayed-model" assert mock_page.evaluate.call_count == 0 @pytest.mark.asyncio async def test_set_model_from_display_with_storage(mock_page, mock_server): """Test setting model from page display and updating storage.""" mock_element = AsyncMock() mock_element.inner_text.return_value = "displayed-model" mock_locator = AsyncMock() mock_locator.first = mock_element mock_page.locator.return_value = mock_locator mock_page.evaluate.return_value = json.dumps({}) with ( patch.dict(sys.modules, {"server": mock_server}), patch( "browser_utils.models.startup._verify_and_apply_ui_state", return_value=True, ), ): await _set_model_from_page_display(mock_page, set_storage=True) assert any( "setItem" in str(args) for args, _ in mock_page.evaluate.call_args_list ) @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_set_model_from_page_display_success(mock_page, mock_server): """Test _set_model_from_page_display success.""" mock_locator = MagicMock() mock_element = AsyncMock() mock_element.inner_text.return_value = "gemini-ultra" type(mock_locator).first = PropertyMock(return_value=mock_element) mock_page.locator.return_value = mock_locator with ( patch.dict("sys.modules", {"server": mock_server}), patch( "browser_utils.models.startup._verify_and_apply_ui_state", return_value=True, ), ): await _set_model_from_page_display(mock_page, set_storage=True) assert mock_server.current_ai_studio_model_id == "gemini-ultra" assert mock_page.evaluate.called @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_set_model_from_page_display_set_storage_defaults(mock_page, mock_server): """Test set_storage=True logic with default keys.""" mock_page.locator.return_value.first.inner_text = AsyncMock( return_value="gemini-pro" ) mock_page.evaluate.return_value = None # No existing prefs with ( patch.dict(sys.modules, {"server": mock_server}), patch( "browser_utils.models.startup._verify_and_apply_ui_state", return_value=True, ), ): await _set_model_from_page_display(mock_page, set_storage=True) args = mock_page.evaluate.call_args[0] assert "localStorage.setItem" in args[0] prefs = json.loads(args[1]) assert prefs["isAdvancedOpen"] is True assert prefs["promptModel"] == "models/gemini-pro" assert prefs["getCodeLanguage"] == "Node.js" @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_set_model_from_page_display_same_id(mock_page): """Test when displayed ID matches current server ID.""" mock_state = MagicMock() mock_state.current_ai_studio_model_id = "gemini-pro" mock_state.parsed_model_list = [] mock_event = asyncio.Event() mock_event.set() # Already set mock_state.model_list_fetch_event = mock_event mock_page.locator.return_value.first.inner_text = AsyncMock( return_value="gemini-pro" ) with ( patch("api_utils.server_state.state", mock_state), patch("browser_utils.models.startup.logger") as mock_logger, ): await _set_model_from_page_display(mock_page) # Model ID should not have changed assert mock_state.current_ai_studio_model_id == "gemini-pro" # Implementation doesn't log when unchanged, so just verify debug was called for reading debug_calls = [call[0][0] for call in mock_logger.debug.call_args_list] assert any("gemini-pro" in msg for msg in debug_calls) # === Section 6: Switch Model Tests === @pytest.mark.asyncio async def test_switch_model_recovery_logic(mock_page): """Test switch model failure triggering recovery logic.""" req_id = "test_req" model_id = "new-model" mock_page.evaluate.return_value = json.dumps({"promptModel": "models/old-model"}) mock_page.url = "http://example.com" mock_locator = AsyncMock() mock_locator.first.inner_text.return_value = "old-model" mock_page.locator.return_value = mock_locator with ( patch( "server.parsed_model_list", [{"id": "old-model", "display_name": "Old Model"}], ), patch( "browser_utils.models.switcher._verify_and_apply_ui_state", return_value=True, ), patch("browser_utils.models.switcher.expect_async") as mock_expect, ): mock_expect.return_value.to_be_visible = AsyncMock() result = await switch_ai_studio_model(mock_page, model_id, req_id) found_target = False for call in mock_page.evaluate.call_args_list: args, _ = call if len(args) > 1 and "models/new-model" in str(args[1]): found_target = True break assert found_target assert result is False @pytest.mark.asyncio async def test_switch_model_json_error_original(mock_page): """Test switch model when original prefs are invalid JSON.""" mock_page.evaluate.return_value = "invalid json" mock_locator = AsyncMock() mock_locator.first.inner_text.return_value = "unknown-model" mock_page.locator.return_value = mock_locator with patch( "browser_utils.models.switcher._verify_and_apply_ui_state", return_value=True ): result = await switch_ai_studio_model(mock_page, "new-model", "req_id") assert result is False @pytest.mark.asyncio async def test_switch_model_already_set_nav_needed(mock_page): """Test when model is already set in storage but URL is wrong.""" mock_page.evaluate.return_value = json.dumps({"promptModel": "models/target-model"}) mock_page.url = "http://wrong.url" mock_locator = AsyncMock() mock_locator.first.inner_text.return_value = "target-model" mock_incognito = AsyncMock() mock_incognito.get_attribute.return_value = ["ms-button-active"] def locator_side_effect(selector): if "model-name" in selector: return mock_locator if "Temporary chat toggle" in selector: return mock_incognito return AsyncMock() mock_page.locator.side_effect = locator_side_effect mock_expect = MagicMock() mock_expect.to_be_visible = AsyncMock() with ( patch( "browser_utils.models.switcher._verify_and_apply_ui_state", return_value=True, ), patch("browser_utils.models.switcher.expect_async", return_value=mock_expect), ): result = await switch_ai_studio_model(mock_page, "target-model", "req_id") assert result is True mock_page.goto.assert_called() @pytest.mark.asyncio async def test_switch_model_success_flow(mock_page, mock_server): """Test full success flow: not set -> set -> navigate -> verify -> success.""" req_id = "test_req" model_id = "target-model" initial_state = json.dumps({"promptModel": "models/old-model"}) final_state = json.dumps({"promptModel": "models/target-model"}) mock_page.evaluate.side_effect = [ initial_state, # 1. get initial None, # 2. set item None, # 3. set item (compat) final_state, # 4. get final initial_state, # 5. get revert (if failure) None, # 6. set revert ] mock_element = AsyncMock() mock_element.inner_text.return_value = "target-model" mock_locator = AsyncMock() mock_locator.first = mock_element mock_incognito = AsyncMock() mock_incognito.get_attribute.return_value = "ms-button-active custom-class" mock_incognito.wait_for.return_value = None mock_incognito.click.return_value = None def locator_side_effect(selector): if "model-name" in selector: return mock_locator if "Temporary chat toggle" in selector: return mock_incognito return AsyncMock() mock_page.locator.side_effect = locator_side_effect mock_expect = MagicMock() mock_expect.to_be_visible = AsyncMock() mock_server.parsed_model_list = [ {"id": "target-model", "display_name": "Target Model"} ] with ( patch( "browser_utils.models.switcher._verify_and_apply_ui_state", return_value=True, ), patch("browser_utils.models.switcher.expect_async", return_value=mock_expect), patch.dict(sys.modules, {"server": mock_server}), ): result = await switch_ai_studio_model(mock_page, model_id, req_id) assert result is True mock_page.goto.assert_called() assert any("setItem" in str(args) for args, _ in mock_page.evaluate.call_args_list) @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_switch_model_json_error_original_prefs(mock_page): """Test JSONDecodeError when parsing original prefs.""" mock_page.evaluate.side_effect = [ "invalid-json", # 1. original None, # 2. set promptModel None, # 3. set UI state manual json.dumps({"promptModel": "models/target-model"}), # 4. final ] mock_page.url = "https://aistudio.google.com/prompts/new_chat" mock_locator = MagicMock() mock_element = AsyncMock() mock_element.inner_text.return_value = "target-model" type(mock_locator).first = PropertyMock(return_value=mock_element) mock_page.locator.return_value = mock_locator mock_locator.get_attribute = AsyncMock(return_value=["ms-button-active"]) with ( patch( "browser_utils.models.switcher._verify_and_apply_ui_state", return_value=True, ), patch("browser_utils.models.switcher.expect_async") as mock_expect, ): mock_expect.return_value.to_be_visible = AsyncMock() result = await switch_ai_studio_model(mock_page, "target-model", "req1") assert result is True @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_switch_model_json_error_final_prefs(mock_page): """Test JSONDecodeError when parsing final prefs.""" initial_prefs = json.dumps({"promptModel": "models/old-model"}) mock_page.evaluate.side_effect = [ initial_prefs, # 1 None, # 2 None, # 3 "invalid-json", # 4 (final check -> fails) initial_prefs, # 5 (revert: read current LS) None, # 6 (revert: write LS) ] mock_page.url = "https://aistudio.google.com/prompts/new_chat" mock_locator = MagicMock() mock_element = AsyncMock() mock_element.inner_text.return_value = "target-model" type(mock_locator).first = PropertyMock(return_value=mock_element) mock_incognito = MagicMock() mock_incognito.wait_for = AsyncMock() mock_incognito.get_attribute = AsyncMock(return_value=["ms-button-active"]) mock_incognito.click = AsyncMock() def locator_side_effect(selector): if "model-name" in selector: return mock_locator if "Temporary chat toggle" in selector: return mock_incognito return MagicMock() mock_page.locator.side_effect = locator_side_effect with ( patch( "browser_utils.models.switcher._verify_and_apply_ui_state", return_value=True, ), patch("browser_utils.models.switcher.expect_async") as mock_expect, ): mock_expect.return_value.to_be_visible = AsyncMock() result = await switch_ai_studio_model(mock_page, "target-model", "req1") assert result is False @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_switch_model_read_model_name_exception(mock_page): """Test exception when reading displayed model name.""" initial_prefs = json.dumps({"promptModel": "models/old-model"}) mock_page.evaluate.side_effect = [ initial_prefs, None, None, json.dumps({"promptModel": "models/target-model"}), None, # Revert to original ] mock_page.url = "https://aistudio.google.com/prompts/new_chat" mock_locator = MagicMock() mock_element = AsyncMock() mock_element.inner_text.side_effect = Exception("Read Error") type(mock_locator).first = PropertyMock(return_value=mock_element) def locator_side_effect(selector): if "model-name" in selector: return mock_locator return MagicMock() mock_page.locator.side_effect = locator_side_effect with ( patch( "browser_utils.models.switcher._verify_and_apply_ui_state", return_value=True, ), patch("browser_utils.models.switcher.expect_async") as mock_expect, ): mock_expect.return_value.to_be_visible = AsyncMock() result = await switch_ai_studio_model(mock_page, "target-model", "req1") assert result is False @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_switch_model_incognito_retry(mock_page): """Test retrying to enable incognito mode.""" initial_prefs = json.dumps({"promptModel": "models/old-model"}) mock_page.evaluate.side_effect = [ initial_prefs, None, None, json.dumps({"promptModel": "models/target-model"}), ] mock_page.url = "https://aistudio.google.com/prompts/new_chat" mock_locator = MagicMock() mock_element = AsyncMock() mock_element.inner_text.return_value = "target-model" type(mock_locator).first = PropertyMock(return_value=mock_element) mock_incognito = MagicMock() mock_incognito.wait_for = AsyncMock() mock_incognito.get_attribute = AsyncMock(side_effect=[[], ["ms-button-active"]]) mock_incognito.click = AsyncMock() def locator_side_effect(selector): if "model-name" in selector: return mock_locator if "Temporary chat toggle" in selector: return mock_incognito return MagicMock() mock_page.locator.side_effect = locator_side_effect with ( patch( "browser_utils.models.switcher._verify_and_apply_ui_state", return_value=True, ), patch("browser_utils.models.switcher.expect_async") as mock_expect, ): mock_expect.return_value.to_be_visible = AsyncMock() result = await switch_ai_studio_model(mock_page, "target-model", "req1") assert result is True assert mock_incognito.click.called @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_switch_model_revert_cant_read_display(mock_page): """Test revert flow when display name cannot be read.""" initial_prefs = json.dumps({"promptModel": "models/original-model"}) mock_page.evaluate.side_effect = [ initial_prefs, # 1 None, # 2 None, # 3 json.dumps({"promptModel": "models/target-model"}), # 4 None, # 5 ] mock_page.url = "https://aistudio.google.com/prompts/new_chat" mock_locator = MagicMock() mock_element = AsyncMock() mock_element.inner_text.return_value = "wrong-model" type(mock_locator).first = PropertyMock(return_value=mock_element) mock_revert_locator = MagicMock() mock_revert_element = AsyncMock() mock_revert_element.inner_text.side_effect = Exception("Revert Read Error") type(mock_revert_locator).first = PropertyMock(return_value=mock_revert_element) locators = iter([mock_locator, mock_revert_locator]) def locator_side_effect(selector): if "model-name" in selector: return next(locators) return MagicMock() mock_page.locator.side_effect = locator_side_effect with ( patch( "browser_utils.models.switcher._verify_and_apply_ui_state", return_value=True, ), patch("browser_utils.models.switcher.expect_async") as mock_expect, ): mock_expect.return_value.to_be_visible = AsyncMock() result = await switch_ai_studio_model(mock_page, "target-model", "req1") assert result is False set_calls = [ args for args, _ in mock_page.evaluate.call_args_list if "setItem" in str(args) ] # Should have set the target model, but no revert last_set_arg = set_calls[-1][1] assert "target-model" in last_set_arg @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_switch_model_general_exception(mock_page): """Test general exception handling in switch_ai_studio_model.""" mock_page.evaluate.side_effect = Exception("Catastrophic Failure") result = await switch_ai_studio_model(mock_page, "target-model", "req1") assert result is False @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_switch_ai_studio_model_json_error_logging(mock_page): """Test handling of JSONDecodeError in original prefs with logging.""" mock_page.evaluate.side_effect = [ "invalid-json", # original_prefs_str None, # setItem 1 None, # setItem 2 None, # final_prefs_str (check) None, # check for revert (current_ls_content_str) ] with ( patch("browser_utils.models.switcher.logger") as mock_logger, patch( "browser_utils.models.switcher._verify_and_apply_ui_state", return_value=True, ), patch("browser_utils.models.switcher.expect_async") as mock_expect, ): mock_expect.return_value.to_be_visible = AsyncMock() mock_page.locator.return_value.first.inner_text = AsyncMock( return_value="Different Model" ) await switch_ai_studio_model(mock_page, "gemini-pro", "req1") warnings = [call.args[0] for call in mock_logger.warning.call_args_list] assert any( "failed to parse original aistudiouserpreference json string" in w.lower() for w in warnings ) @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_switch_ai_studio_model_ui_state_fail(mock_page): """Test warning when UI state verification fails but process continues.""" mock_page.evaluate.return_value = json.dumps({"promptModel": "models/old-model"}) with ( patch("browser_utils.models.switcher.logger") as mock_logger, patch( "browser_utils.models.switcher._verify_and_apply_ui_state", return_value=False, ), patch("browser_utils.models.switcher.expect_async") as mock_expect, ): mock_expect.return_value.to_be_visible = AsyncMock() mock_page.locator.return_value.first.inner_text = AsyncMock( return_value="gemini-pro" ) await switch_ai_studio_model(mock_page, "gemini-pro", "req1") warnings = [call.args[0] for call in mock_logger.warning.call_args_list] assert any( "UI state setting failed, but continuing model switching flow" in str(w) for w in warnings ) @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_switch_ai_studio_model_final_storage_mismatch(mock_page): """Test error when final storage does not match target model.""" mock_page.evaluate.side_effect = [ json.dumps({"promptModel": "models/old-model"}), # 1. Get original None, # 2. Set item 1 None, # 3. Set item 2 json.dumps({"promptModel": "models/old-model"}), # 4. Get final (still old) None, # 5. Set item (revert) json.dumps({"promptModel": "models/old-model"}), # 6. Get item (revert logic) None, # 7. Set item (revert) ] with ( patch("browser_utils.models.switcher.logger") as mock_logger, patch( "browser_utils.models.switcher._verify_and_apply_ui_state", return_value=True, ), patch("browser_utils.models.switcher.expect_async") as mock_expect, ): mock_expect.return_value.to_be_visible = AsyncMock() mock_page.locator.return_value.first.inner_text = AsyncMock( return_value="gemini-pro" ) await switch_ai_studio_model(mock_page, "gemini-pro", "req1") errors = [call.args[0] for call in mock_logger.error.call_args_list] assert any("AI Studio did not accept model change" in e for e in errors) # === Section 7: Revert Logic Tests === @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_switch_model_revert_success(mock_page, mock_server): """Test revert logic when switch fails but revert works by reading displayed model.""" req_id = "test_revert_success" model_id = "target-model" with patch.dict(sys.modules, {"server": mock_server}): initial_state = json.dumps({"promptModel": "models/old-model"}) final_state = json.dumps({"promptModel": "models/target-model"}) mock_page.evaluate.side_effect = [ initial_state, # 1. get initial None, # 2. set item None, # 3. set item (compat) final_state, # 4. consumed by revert logic (get storage) initial_state, # 5. consumed by revert logic (set revert storage) None, # 6. set revert storage (extra safety) ] + [None] * 20 mock_element = AsyncMock() mock_element.inner_text.return_value = "old-model" mock_locator = MagicMock() type(mock_locator).first = PropertyMock(return_value=mock_element) mock_page.locator.return_value = mock_locator mock_server.parsed_model_list = [ {"id": "target-model", "display_name": "Target Model"}, {"id": "old-model", "display_name": "Old Model"}, ] mock_expect = MagicMock() mock_expect.to_be_visible = AsyncMock() with ( patch( "browser_utils.models.switcher._verify_and_apply_ui_state", new_callable=AsyncMock, side_effect=[False, True, True, True], ) as mock_verify, patch( "browser_utils.models.switcher.expect_async", return_value=mock_expect ), patch( "browser_utils.operations.save_error_snapshot", new_callable=AsyncMock ), ): result = await switch_ai_studio_model(mock_page, model_id, req_id) assert result is False set_calls = [ args for args, _ in mock_page.evaluate.call_args_list if "setItem" in str(args) ] assert len(set_calls) == 2 last_set_arg = set_calls[-1][1] assert "models/target-model" in last_set_arg assert mock_verify.call_count == 2 @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_switch_model_revert_failure_fallback(mock_page, mock_server): """Test revert logic when reading displayed model fails, falls back to original storage.""" req_id = "test_revert_fallback" model_id = "target-model" initial_state = json.dumps({"promptModel": "models/original-model"}) mock_page.evaluate.side_effect = [ initial_state, # 1. get initial None, # 2. set item None, # 3. set item (compat) None, # 4. set fallback storage ] + [None] * 10 mock_locator = MagicMock() mock_element = AsyncMock() mock_element.inner_text.side_effect = Exception("Locator fail") type(mock_locator).first = PropertyMock(return_value=mock_element) mock_page.locator.return_value = mock_locator mock_server.parsed_model_list = [ {"id": "target-model", "display_name": "Target Model"} ] mock_expect = MagicMock() mock_expect.to_be_visible = AsyncMock() with ( patch( "browser_utils.models.switcher._verify_and_apply_ui_state", new_callable=AsyncMock, return_value=False, ), patch("browser_utils.models.switcher.expect_async", return_value=mock_expect), ): result = await switch_ai_studio_model(mock_page, model_id, req_id) assert result is False set_calls = [ args for args, _ in mock_page.evaluate.call_args_list if "setItem" in str(args) ] last_set_arg = set_calls[-1][1] assert "models/target-model" in last_set_arg @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_switch_model_revert_blind_trust(mock_page, mock_server): """Test revert logic uses displayed model name as ID blindly (current behavior).""" req_id = "test_revert_blind" model_id = "target-model" initial_state = json.dumps({"promptModel": "models/original-model"}) mock_page.evaluate.side_effect = [ initial_state, # 1. get initial None, # 2. set item None, # 3. set item (compat) initial_state, # 4. get revert storage None, # 5. set revert storage ] + [None] * 10 mock_element = AsyncMock() mock_element.inner_text.return_value = "Unknown Model" mock_locator = MagicMock() type(mock_locator).first = PropertyMock(return_value=mock_element) mock_page.locator.return_value = mock_locator mock_server.parsed_model_list = [ {"id": "target-model", "display_name": "Target Model"} ] mock_expect = MagicMock() mock_expect.to_be_visible = AsyncMock() with ( patch( "browser_utils.models.switcher._verify_and_apply_ui_state", new_callable=AsyncMock, side_effect=[False, True, True, True], ), patch("browser_utils.models.switcher.expect_async", return_value=mock_expect), patch("browser_utils.operations.save_error_snapshot", new_callable=AsyncMock), ): result = await switch_ai_studio_model(mock_page, model_id, req_id) assert result is False set_calls = [ args for args, _ in mock_page.evaluate.call_args_list if "setItem" in str(args) ] last_set_arg = set_calls[-1][1] assert "models/target-model" in last_set_arg