tests/browser_utils/test_model_management_coverage.py

# Focused coverage tests for browser_utils/model_management.py
# Targets specific missing lines to achieve >80% coverage

import asyncio
import json
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

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():
    """Simple mock page."""
    page = AsyncMock()
    page.locator = MagicMock()
    page.evaluate = AsyncMock(return_value=None)
    page.goto = AsyncMock()
    page.url = "https://aistudio.google.com/prompts/new_chat"
    return page


# ===== _verify_ui_state_settings Coverage =====


@pytest.mark.asyncio
async def test_verify_ui_missing_storage(mock_page):
    """Lines 49-57: localStorage missing."""
    mock_page.evaluate.return_value = None

    result = await _verify_ui_state_settings(mock_page, "req1")

    assert result["exists"] is False
    assert result["error"] == "localStorage not found"
    assert result["needsUpdate"] is True


@pytest.mark.asyncio
async def test_verify_ui_json_error(mock_page):
    """Lines 82-90: JSONDecodeError handling."""
    mock_page.evaluate.return_value = "invalid json"

    result = await _verify_ui_state_settings(mock_page, "req1")

    assert result["exists"] is False
    assert "JSON parse failed" in result["error"]


@pytest.mark.asyncio
async def test_verify_ui_general_exception(mock_page):
    """Lines 94-102: General exception handling."""
    mock_page.evaluate.side_effect = Exception("Page error")

    result = await _verify_ui_state_settings(mock_page, "req1")

    assert result["exists"] is False
    assert "Verification failed" in result["error"]


# ===== _force_ui_state_settings Coverage =====


@pytest.mark.asyncio
async def test_force_ui_no_update_needed(mock_page):
    """Lines 122-124: Early return when no update needed."""
    with patch(
        "browser_utils.models.ui_state._verify_ui_state_settings",
        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
async def test_force_ui_verify_fail(mock_page):
    """Lines 147-149: Verification fails after setting."""
    with patch(
        "browser_utils.models.ui_state._verify_ui_state_settings",
        side_effect=[{"needsUpdate": True, "prefs": {}}, {"needsUpdate": True}],
    ):
        result = await _force_ui_state_settings(mock_page, "req1")

        assert result is False


# ===== _force_ui_state_with_retry Coverage =====


@pytest.mark.asyncio
async def test_retry_success_first_attempt(mock_page):
    """Lines 180-182: Success on first attempt."""
    with patch(
        "browser_utils.models.ui_state._force_ui_state_settings", return_value=True
    ):
        result = await _force_ui_state_with_retry(mock_page, max_retries=3)

        assert result is True


@pytest.mark.asyncio
async def test_retry_fail_all(mock_page):
    """Lines 184-189: All retries fail."""
    with patch(
        "browser_utils.models.ui_state._force_ui_state_settings", return_value=False
    ):
        result = await _force_ui_state_with_retry(
            mock_page, max_retries=2, retry_delay=0.01
        )

        assert result is False


# ===== _verify_and_apply_ui_state Coverage =====


@pytest.mark.asyncio
async def test_apply_ui_needs_update(mock_page):
    """Lines 214-216: Needs update path."""
    with (
        patch(
            "browser_utils.models.ui_state._verify_ui_state_settings",
            return_value={
                "exists": True,
                "needsUpdate": True,
                "isAdvancedOpen": False,
                "areToolsOpen": False,
            },
        ),
        patch(
            "browser_utils.models.ui_state._force_ui_state_with_retry",
            return_value=True,
        ),
    ):
        result = await _verify_and_apply_ui_state(mock_page, "req1")

        assert result is True


@pytest.mark.asyncio
async def test_apply_ui_already_ok(mock_page):
    """Lines 217-219: No update needed."""
    with patch(
        "browser_utils.models.ui_state._verify_ui_state_settings",
        return_value={
            "exists": True,
            "needsUpdate": False,
            "isAdvancedOpen": True,
            "areToolsOpen": True,
        },
    ):
        result = await _verify_and_apply_ui_state(mock_page, "req1")

        assert result is True


# ===== switch_ai_studio_model Coverage =====


@pytest.mark.asyncio
async def test_switch_model_json_error_original(mock_page):
    """Lines 246-248: JSONDecodeError on original prefs."""
    mock_page.evaluate.side_effect = ["invalid json", None, None, None]

    mock_locator = MagicMock()
    mock_locator.first.inner_text = AsyncMock(return_value="gemini-pro")
    mock_page.locator.return_value = mock_locator

    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, "gemini-pro", "req1")

        # Should handle error and continue
        assert result in [True, False]  # May succeed or fail depending on verification


@pytest.mark.asyncio
async def test_switch_model_already_set_wrong_url(mock_page):
    """Lines 256-269: Model already set but URL wrong."""
    prefs = json.dumps({"promptModel": "models/gemini-pro"})
    mock_page.evaluate.return_value = prefs
    mock_page.url = "https://wrong.url"

    with 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, "gemini-pro", "req1")

        assert result is True
        mock_page.goto.assert_called_once()


@pytest.mark.asyncio
async def test_switch_model_ui_state_fail_warning(mock_page):
    """Lines 283-284: UI state fails but continues."""
    prefs = json.dumps({"promptModel": "models/old"})
    mock_page.evaluate.side_effect = [
        prefs,
        None,
        None,
        json.dumps({"promptModel": "models/new"}),
    ]

    mock_locator = MagicMock()
    mock_locator.first.inner_text = AsyncMock(return_value="new")
    mock_page.locator.return_value = mock_locator

    with (
        patch(
            "browser_utils.models.switcher._verify_and_apply_ui_state",
            return_value=False,
        ),
        patch("browser_utils.models.switcher.expect_async") as mock_expect,
        patch("browser_utils.models.switcher.logger") as mock_logger,
    ):
        mock_expect.return_value.to_be_visible = AsyncMock()

        await switch_ai_studio_model(mock_page, "new", "req1")

        # Verify warning logged
        warnings = [call.args[0] for call in mock_logger.warning.call_args_list]
        assert any("UI state setting failed" in str(w) for w in warnings)


@pytest.mark.asyncio
async def test_switch_model_final_ui_fail(mock_page):
    """Lines 303-307: Final UI state verification."""
    prefs = json.dumps({"promptModel": "models/old"})
    mock_page.evaluate.side_effect = [
        prefs,
        None,
        None,
        json.dumps({"promptModel": "models/new"}),
    ]

    mock_locator = MagicMock()
    mock_locator.first.inner_text = AsyncMock(return_value="new")
    mock_page.locator.return_value = mock_locator

    ui_calls = [True, False]  # Initial succeeds, final fails

    with (
        patch(
            "browser_utils.models.switcher._verify_and_apply_ui_state",
            side_effect=lambda *args: ui_calls.pop(0),
        ),
        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, "new", "req1")

        assert result is True  # Still succeeds despite warning


@pytest.mark.asyncio
async def test_switch_model_final_prefs_json_error(mock_page):
    """Lines 317-318: JSONDecodeError on final prefs."""
    prefs = json.dumps({"promptModel": "models/old"})
    mock_page.evaluate.side_effect = [prefs, None, None, "invalid json"]

    with (
        patch(
            "browser_utils.models.switcher._verify_and_apply_ui_state",
            return_value=True,
        ),
        patch("browser_utils.models.switcher.expect_async") as mock_expect,
        patch("browser_utils.models.switcher.logger"),
    ):
        mock_expect.return_value.to_be_visible = AsyncMock()

        result = await switch_ai_studio_model(mock_page, "new", "req1")

        assert result is False  # Fails due to verification failure


@pytest.mark.asyncio
async def test_switch_model_display_read_error(mock_page):
    """Lines 360-366: Exception reading displayed model."""
    prefs = json.dumps({"promptModel": "models/old"})
    mock_page.evaluate.side_effect = [
        prefs,
        None,
        None,
        json.dumps({"promptModel": "models/new"}),
    ]

    mock_locator = MagicMock()
    mock_locator.first.inner_text = AsyncMock(side_effect=Exception("Read failed"))
    mock_page.locator.return_value = mock_locator

    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, "new", "req1")

        assert result is False


@pytest.mark.asyncio
async def test_switch_model_incognito_active(mock_page):
    """Lines 383-384: Incognito already active."""
    prefs = json.dumps({"promptModel": "models/old"})
    mock_page.evaluate.side_effect = [
        prefs,
        None,
        None,
        json.dumps({"promptModel": "models/new"}),
    ]

    mock_model_loc = MagicMock()
    mock_model_loc.first.inner_text = AsyncMock(return_value="new")

    mock_incognito = MagicMock()
    mock_incognito.wait_for = AsyncMock()
    mock_incognito.get_attribute = AsyncMock(return_value="ms-button-active")

    def loc_side_effect(sel):
        if "model-name" in sel:
            return mock_model_loc
        if "Temporary" in sel:
            return mock_incognito
        return MagicMock()

    mock_page.locator.side_effect = loc_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, "new", "req1")

        assert result is True


@pytest.mark.asyncio
async def test_switch_model_incognito_exception(mock_page):
    """Lines 400-403: Incognito toggle fails."""
    prefs = json.dumps({"promptModel": "models/old"})
    mock_page.evaluate.side_effect = [
        prefs,
        None,
        None,
        json.dumps({"promptModel": "models/new"}),
    ]

    mock_model_loc = MagicMock()
    mock_model_loc.first.inner_text = AsyncMock(return_value="new")

    mock_incognito = MagicMock()
    mock_incognito.wait_for = AsyncMock(side_effect=Exception("Button not found"))

    def loc_side_effect(sel):
        if "model-name" in sel:
            return mock_model_loc
        if "Temporary" in sel:
            return mock_incognito
        return MagicMock()

    mock_page.locator.side_effect = loc_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, "new", "req1")

        assert result is True  # Succeeds despite incognito failure


# ===== load_excluded_models Coverage =====


@pytest.mark.asyncio
async def test_load_excluded_file_exists():
    """Lines 597-609: Successful load."""
    from api_utils.server_state import state

    state.excluded_model_ids = set()

    with (
        patch("os.path.exists", return_value=True),
        patch("builtins.open", new_callable=MagicMock) as mock_open,
    ):
        mock_file = MagicMock()
        mock_file.__enter__.return_value = ["model-1\n", "model-2\n"]
        mock_open.return_value = mock_file

        load_excluded_models("test.txt")

        assert "model-1" in state.excluded_model_ids
        assert "model-2" in state.excluded_model_ids


@pytest.mark.asyncio
async def test_load_excluded_empty_file():
    """Lines 606-609: Empty file."""
    from api_utils.server_state import state

    state.excluded_model_ids = set()

    with (
        patch("os.path.exists", return_value=True),
        patch("builtins.open", new_callable=MagicMock) as mock_open,
    ):
        mock_file = MagicMock()
        mock_file.__enter__.return_value = []
        mock_open.return_value = mock_file

        load_excluded_models("empty.txt")

        assert len(state.excluded_model_ids) == 0


@pytest.mark.asyncio
async def test_load_excluded_file_not_found():
    """Lines 610-611: File not found."""
    from api_utils.server_state import state

    state.excluded_model_ids = set()

    with patch("os.path.exists", return_value=False):
        load_excluded_models("nonexistent.txt")

        assert len(state.excluded_model_ids) == 0


@pytest.mark.asyncio
async def test_load_excluded_exception():
    """Lines 612-613: Exception during load."""
    from api_utils.server_state import state

    state.excluded_model_ids = set()

    with patch("os.path.exists", side_effect=Exception("Disk error")):
        load_excluded_models("error.txt")

        # Should not crash, just log error
        assert len(state.excluded_model_ids) == 0


# ===== _handle_initial_model_state_and_storage Coverage =====


@pytest.mark.asyncio
async def test_handle_initial_missing_storage(mock_page):
    """Lines 632-635: Missing localStorage."""

    mock_page.evaluate.return_value = None
    mock_page.url = "https://test.url"

    with (
        patch("browser_utils.models.startup._set_model_from_page_display"),
        patch(
            "browser_utils.models.startup._verify_and_apply_ui_state",
            return_value=True,
        ),
        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 trigger reload flow
        assert mock_page.goto.called


@pytest.mark.asyncio
async def test_handle_initial_json_error(mock_page):
    """Lines 664-669: JSONDecodeError."""
    mock_page.evaluate.return_value = "invalid json"
    mock_page.url = "https://test.url"

    with (
        patch("browser_utils.models.startup._set_model_from_page_display"),
        patch(
            "browser_utils.models.startup._verify_and_apply_ui_state",
            return_value=True,
        ),
        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 trigger reload
        assert mock_page.goto.called


@pytest.mark.asyncio
async def test_handle_initial_exception_fallback(mock_page):
    """Lines 738-753: Exception with fallback."""
    mock_page.evaluate.side_effect = Exception("Critical error")

    with patch(
        "browser_utils.models.startup._set_model_from_page_display"
    ) as mock_fallback:
        await _handle_initial_model_state_and_storage(mock_page)

        # Should call fallback
        mock_fallback.assert_called_once()
        assert mock_fallback.call_args[1]["set_storage"] is False


# ===== _set_model_from_page_display Coverage =====


@pytest.mark.asyncio
async def test_set_model_display_basic(mock_page):
    """Lines 765-795: Basic display read."""
    from api_utils.server_state import state

    state.current_ai_studio_model_id = "old"

    mock_loc = MagicMock()
    mock_loc.first.inner_text = AsyncMock(return_value="new-model")
    mock_page.locator.return_value = mock_loc

    await _set_model_from_page_display(mock_page, set_storage=False)

    assert state.current_ai_studio_model_id == "new-model"


@pytest.mark.asyncio
async def test_set_model_display_with_storage(mock_page):
    """Lines 797-860: Storage update path."""
    from api_utils.server_state import state

    # Setup event mock
    mock_event = AsyncMock(spec=asyncio.Event)
    mock_event.is_set = MagicMock(return_value=True)
    state.model_list_fetch_event = mock_event

    mock_loc = MagicMock()
    mock_loc.first.inner_text = AsyncMock(return_value="model-id")
    mock_page.locator.return_value = mock_loc

    # First call returns existing prefs, second call for setItem
    mock_page.evaluate.side_effect = [json.dumps({}), None]

    with patch(
        "browser_utils.models.startup._verify_and_apply_ui_state", return_value=True
    ):
        await _set_model_from_page_display(mock_page, set_storage=True)

        # Verify setItem called
        calls = [str(c) for c in mock_page.evaluate.call_args_list]
        assert any("setItem" in c for c in calls)


@pytest.mark.asyncio
async def test_set_model_display_json_error_storage(mock_page):
    """Lines 807-811: JSONDecodeError on existing prefs."""
    from api_utils.server_state import state

    # Setup event mock
    mock_event = AsyncMock(spec=asyncio.Event)
    mock_event.is_set = MagicMock(return_value=True)
    state.model_list_fetch_event = mock_event

    mock_loc = MagicMock()
    mock_loc.first.inner_text = AsyncMock(return_value="model")
    mock_page.locator.return_value = mock_loc

    # First call returns invalid JSON, second call for setItem
    mock_page.evaluate.side_effect = ["invalid json", None]

    with patch(
        "browser_utils.models.startup._verify_and_apply_ui_state", return_value=True
    ):
        await _set_model_from_page_display(mock_page, set_storage=True)

        # Should handle error and create new prefs
        assert mock_page.evaluate.call_count >= 2


@pytest.mark.asyncio
async def test_set_model_display_exception(mock_page):
    """Lines 861-864: Exception handling."""
    mock_loc = MagicMock()
    mock_loc.first.inner_text = AsyncMock(side_effect=Exception("Read failed"))
    mock_page.locator.return_value = mock_loc

    # Should not raise
    await _set_model_from_page_display(mock_page, set_storage=False)


# ===== Updated Switch Model Coverage (No revert logic anymore) =====


@pytest.mark.asyncio
async def test_switch_model_validation_fail_storage(mock_page):
    """Cover failure path when storage does not match target."""
    original_prefs = json.dumps({"promptModel": "models/old"})
    mock_page.evaluate.side_effect = [
        original_prefs,  # get original
        None,  # set target
        None,  # set target (compat)
        original_prefs,  # get final (mismatch!)
    ]

    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, "new", "req1")
        assert result is False


@pytest.mark.asyncio
async def test_switch_model_validation_fail_display(mock_page):
    """Cover failure path when page display does not match target."""
    prefs = json.dumps({"promptModel": "models/new"})
    mock_page.evaluate.side_effect = [
        json.dumps({"promptModel": "models/old"}),  # original
        None,  # set
        None,  # set compat
        prefs,  # final check ok
    ]

    mock_locator = MagicMock()
    mock_locator.first.inner_text = AsyncMock(return_value="wrong-display")
    mock_page.locator.return_value = mock_locator

    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, "new", "req1")
        assert result is False


# ===== Additional Coverage for _handle_initial_model_state_and_storage =====


@pytest.mark.asyncio
async def test_handle_initial_reload_retry_all_fail(mock_page):
    """Lines 707-725: All reload retries fail."""
    from api_utils.server_state import state

    state.model_list_fetch_event = AsyncMock(spec=asyncio.Event)
    state.model_list_fetch_event.is_set = MagicMock(return_value=True)

    mock_page.evaluate.return_value = None  # Missing localStorage
    mock_page.url = "https://test.url"
    mock_page.goto.side_effect = Exception("Reload always fails")

    with (
        patch("browser_utils.models.startup._set_model_from_page_display"),
        patch(
            "browser_utils.models.startup._verify_and_apply_ui_state",
            return_value=False,
        ),
        patch("browser_utils.models.startup.expect_async") as mock_expect,
        patch("browser_utils.operations.save_error_snapshot", new_callable=AsyncMock),
        patch(
            "browser_utils.models.startup.asyncio.sleep", new_callable=AsyncMock
        ),  # Skip retry delays
    ):
        mock_expect.return_value.to_be_visible = AsyncMock()

        # Should not raise despite all retries failing
        await _handle_initial_model_state_and_storage(mock_page)

        # Should have tried 3 times
        assert mock_page.goto.call_count == 3


@pytest.mark.asyncio
async def test_handle_initial_valid_state_no_reload(mock_page):
    """Lines 734-737: Valid state, no reload needed."""
    from api_utils.server_state import state

    state.current_ai_studio_model_id = None

    prefs = json.dumps(
        {
            "promptModel": "models/valid-model",
            "isAdvancedOpen": True,
            "areToolsOpen": True,
        }
    )
    mock_page.evaluate.return_value = prefs

    with patch(
        "browser_utils.models.startup._verify_ui_state_settings",
        return_value={"needsUpdate": False},
    ):
        await _handle_initial_model_state_and_storage(mock_page)

        # Should not call goto since state is valid
        mock_page.goto.assert_not_called()
        assert state.current_ai_studio_model_id == "valid-model"


@pytest.mark.asyncio
async def test_set_model_display_wait_for_event_timeout(mock_page):
    """Lines 776-781: Wait for model list event timeout."""
    from api_utils.server_state import state

    mock_event = AsyncMock(spec=asyncio.Event)
    mock_event.is_set = MagicMock(return_value=False)
    mock_event.wait = AsyncMock(side_effect=asyncio.TimeoutError())
    state.model_list_fetch_event = mock_event

    mock_loc = MagicMock()
    mock_loc.first.inner_text = AsyncMock(return_value="model-from-page")
    mock_page.locator.return_value = mock_loc

    await _set_model_from_page_display(mock_page, set_storage=False)

    # Should handle timeout gracefully
    assert state.current_ai_studio_model_id == "model-from-page"


@pytest.mark.asyncio
async def test_set_model_display_ui_state_fail_fallback(mock_page):
    """Lines 816-824: UI state setting fails, use traditional method."""
    from api_utils.server_state import state

    mock_event = AsyncMock(spec=asyncio.Event)
    mock_event.is_set = MagicMock(return_value=True)
    state.model_list_fetch_event = mock_event

    mock_loc = MagicMock()
    mock_loc.first.inner_text = AsyncMock(return_value="test-model")
    mock_page.locator.return_value = mock_loc

    mock_page.evaluate.side_effect = [json.dumps({}), None]

    with patch(
        "browser_utils.models.startup._verify_and_apply_ui_state", return_value=False
    ):
        await _set_model_from_page_display(mock_page, set_storage=True)

        # Should still call setItem with traditional method
        assert mock_page.evaluate.call_count >= 2


@pytest.mark.asyncio
async def test_set_model_display_no_model_id_found(mock_page):
    """Lines 832-835: No model ID found from display."""
    from api_utils.server_state import state

    mock_event = AsyncMock(spec=asyncio.Event)
    mock_event.is_set = MagicMock(return_value=True)
    state.model_list_fetch_event = mock_event

    mock_loc = MagicMock()
    mock_loc.first.inner_text = AsyncMock(return_value="unknown-display")
    mock_page.locator.return_value = mock_loc

    # Return empty prefs without promptModel
    mock_page.evaluate.side_effect = [json.dumps({}), None]

    with patch(
        "browser_utils.models.startup._verify_and_apply_ui_state", return_value=True
    ):
        await _set_model_from_page_display(mock_page, set_storage=True)

        # Should handle missing model ID gracefully
        assert state.current_ai_studio_model_id == "unknown-display"


# ===== Additional Edge Cases for 80% Coverage =====


@pytest.mark.asyncio
async def test_verify_ui_cancellederror(mock_page):
    """Lines 92-93: CancelledError propagation in verify."""
    mock_page.evaluate.side_effect = asyncio.CancelledError()

    with pytest.raises(asyncio.CancelledError):
        await _verify_ui_state_settings(mock_page, "req1")


@pytest.mark.asyncio
async def test_force_ui_cancellederror(mock_page):
    """Lines 151-152: CancelledError in force_ui_state_settings."""
    with patch(
        "browser_utils.models.ui_state._verify_ui_state_settings",
        side_effect=asyncio.CancelledError(),
    ):
        with pytest.raises(asyncio.CancelledError):
            await _force_ui_state_settings(mock_page, "req1")


@pytest.mark.asyncio
async def test_verify_apply_cancellederror(mock_page):
    """Lines 221-222: CancelledError in verify_and_apply."""
    with patch(
        "browser_utils.models.ui_state._verify_ui_state_settings",
        side_effect=asyncio.CancelledError(),
    ):
        with pytest.raises(asyncio.CancelledError):
            await _verify_and_apply_ui_state(mock_page, "req1")


@pytest.mark.asyncio
async def test_switch_model_cancellederror(mock_page):
    """Lines 556-557: CancelledError in switch_ai_studio_model."""
    mock_page.evaluate.side_effect = asyncio.CancelledError()

    with pytest.raises(asyncio.CancelledError):
        await switch_ai_studio_model(mock_page, "new-model", "req1")


@pytest.mark.asyncio
async def test_switch_model_incognito_cancellederror(mock_page):
    """Lines 400-401: CancelledError during incognito toggle."""
    prefs = json.dumps({"promptModel": "models/old"})
    mock_page.evaluate.side_effect = [
        prefs,
        None,
        None,
        json.dumps({"promptModel": "models/new"}),
    ]

    mock_model_loc = MagicMock()
    mock_model_loc.first.inner_text = AsyncMock(return_value="new")

    mock_incognito = MagicMock()
    mock_incognito.wait_for = AsyncMock(side_effect=asyncio.CancelledError())

    def loc_side_effect(sel):
        if "model-name" in sel:
            return mock_model_loc
        if "Temporary" in sel:
            return mock_incognito
        return MagicMock()

    mock_page.locator.side_effect = loc_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()

        with pytest.raises(asyncio.CancelledError):
            await switch_ai_studio_model(mock_page, "new", "req1")


@pytest.mark.asyncio
async def test_switch_model_revert_cancellederror(mock_page):
    """Lines 360-361: CancelledError during display name read."""
    # We trigger CancelledError when reading the displayed model name
    # after navigation, which is part of the validation flow.
    prefs = json.dumps({"promptModel": "models/old"})
    mock_page.evaluate.side_effect = [
        prefs,
        None,
        None,
        json.dumps({"promptModel": "models/new"}),
    ]

    mock_locator = MagicMock()
    mock_locator.first.inner_text = AsyncMock(side_effect=asyncio.CancelledError())
    mock_page.locator.return_value = mock_locator

    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()

        with pytest.raises(asyncio.CancelledError):
            await switch_ai_studio_model(mock_page, "new", "req1")


@pytest.mark.asyncio
async def test_switch_model_exception_recovery_cancellederror(mock_page):
    """Lines 208-209: CancelledError propagation."""
    # We trigger CancelledError during the initial evaluate
    mock_page.evaluate.side_effect = asyncio.CancelledError()

    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()

        with pytest.raises(asyncio.CancelledError):
            await switch_ai_studio_model(mock_page, "new-model", "req1")


@pytest.mark.asyncio
async def test_handle_initial_cancellederror(mock_page):
    """Lines 738-739: CancelledError in handle_initial."""
    mock_page.evaluate.side_effect = asyncio.CancelledError()

    with pytest.raises(asyncio.CancelledError):
        await _handle_initial_model_state_and_storage(mock_page)


@pytest.mark.asyncio
async def test_handle_initial_exception_fallback_cancellederror(mock_page):
    """Lines 750-751: CancelledError in fallback path."""
    mock_page.evaluate.side_effect = Exception("Error")

    with patch(
        "browser_utils.models.startup._set_model_from_page_display",
        side_effect=asyncio.CancelledError(),
    ):
        with pytest.raises(asyncio.CancelledError):
            await _handle_initial_model_state_and_storage(mock_page)


@pytest.mark.asyncio
async def test_set_model_display_cancellederror(mock_page):
    """Lines 861-862: CancelledError in set_model_from_page_display."""
    mock_loc = MagicMock()
    mock_loc.first.inner_text = AsyncMock(side_effect=asyncio.CancelledError())
    mock_page.locator.return_value = mock_loc

    with pytest.raises(asyncio.CancelledError):
        await _set_model_from_page_display(mock_page, set_storage=False)


@pytest.mark.asyncio
async def test_switch_model_display_cancellederror(mock_page):
    """Lines 360-361: CancelledError when reading display model."""
    prefs = json.dumps({"promptModel": "models/old"})
    mock_page.evaluate.side_effect = [
        prefs,
        None,
        None,
        json.dumps({"promptModel": "models/new"}),
    ]

    mock_locator = MagicMock()
    mock_locator.first.inner_text = AsyncMock(side_effect=asyncio.CancelledError())
    mock_page.locator.return_value = mock_locator

    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()

        with pytest.raises(asyncio.CancelledError):
            await switch_ai_studio_model(mock_page, "new", "req1")


@pytest.mark.asyncio
async def test_handle_initial_reload_cancellederror(mock_page):
    """Lines 707-708: CancelledError during page reload."""
    from api_utils.server_state import state

    state.model_list_fetch_event = AsyncMock(spec=asyncio.Event)
    state.model_list_fetch_event.is_set = MagicMock(return_value=True)

    mock_page.evaluate.return_value = None  # Missing localStorage
    mock_page.url = "https://test.url"
    mock_page.goto.side_effect = asyncio.CancelledError()

    with (
        patch("browser_utils.models.startup._set_model_from_page_display"),
        patch(
            "browser_utils.models.startup._verify_and_apply_ui_state",
            return_value=False,
        ),
        patch("browser_utils.models.startup.expect_async") as mock_expect,
    ):
        mock_expect.return_value.to_be_visible = AsyncMock()

        with pytest.raises(asyncio.CancelledError):
            await _handle_initial_model_state_and_storage(mock_page)


@pytest.mark.asyncio
async def test_handle_initial_invalid_promptmodel(mock_page):
    """Lines 646-649: Invalid promptModel in localStorage."""

    # promptModel is empty string (invalid)
    prefs = json.dumps({"promptModel": "   ", "isAdvancedOpen": False})
    mock_page.evaluate.return_value = prefs
    mock_page.url = "https://test.url"

    with (
        patch("browser_utils.models.startup._set_model_from_page_display"),
        patch(
            "browser_utils.models.startup._verify_and_apply_ui_state",
            return_value=True,
        ),
        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 trigger reload due to invalid promptModel
        assert mock_page.goto.called